nucliadb 6.2.1.post2864__py3-none-any.whl → 6.2.1.post2869__py3-none-any.whl

nucliadb/ingest/orm/resource.py

@@ -24,7 +24,7 @@ import logging
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
-from typing import TYPE_CHECKING, Any, AsyncIterator, MutableMapping, Optional, Type
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Type
 
 from nucliadb.common import datamanagers
 from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG
@@ -52,7 +52,6 @@ from nucliadb_protos.resources_pb2 import (
     FieldComputedMetadataWrapper,
     FieldFile,
     FieldID,
-    FieldMetadata,
     FieldQuestionAnswerWrapper,
     FieldText,
     FieldType,
@@ -61,7 +60,6 @@ from nucliadb_protos.resources_pb2 import (
     LinkExtractedData,
     Metadata,
     Paragraph,
-    ParagraphAnnotation,
 )
 from nucliadb_protos.resources_pb2 import Basic as PBBasic
 from nucliadb_protos.resources_pb2 import Conversation as PBConversation
@@ -69,15 +67,6 @@ from nucliadb_protos.resources_pb2 import Extra as PBExtra
 from nucliadb_protos.resources_pb2 import Metadata as PBMetadata
 from nucliadb_protos.resources_pb2 import Origin as PBOrigin
 from nucliadb_protos.resources_pb2 import Relations as PBRelations
-from nucliadb_protos.train_pb2 import (
-    EnabledMetadata,
-    TrainField,
-    TrainMetadata,
-    TrainParagraph,
-    TrainResource,
-    TrainSentence,
-)
-from nucliadb_protos.train_pb2 import Position as TrainPosition
 from nucliadb_protos.utils_pb2 import Relation as PBRelation
 from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils.storages.storage import Storage
@@ -343,36 +332,24 @@ class Resource:
                 )
 
             if self.disable_vectors is False:
-                # XXX: while we don't remove the "default" vectorset concept, we
-                # need to do use None as the default one
-                vo = await field.get_vectors()
-                if vo is not None:
-                    async with datamanagers.with_ro_transaction() as ro_txn:
-                        dimension = await datamanagers.kb.get_matryoshka_vector_dimension(
-                            ro_txn, kbid=self.kb.kbid
-                        )
-                    brain.apply_field_vectors(
-                        field_key,
-                        vo,
-                        matryoshka_vector_dimension=dimension,
-                        replace_field=reindex,
-                    )
-
                 vectorset_configs = []
-                async with datamanagers.with_ro_transaction() as ro_txn:
-                    async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(
-                        ro_txn, kbid=self.kb.kbid
-                    ):
-                        vectorset_configs.append(vectorset_config)
+                async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(
+                    self.txn, kbid=self.kb.kbid
+                ):
+                    vectorset_configs.append(vectorset_config)
+
                 for vectorset_config in vectorset_configs:
-                    vo = await field.get_vectors(vectorset=vectorset_config.vectorset_id)
+                    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
                         brain.apply_field_vectors(
                             field_key,
                             vo,
                             vectorset=vectorset_config.vectorset_id,
-                            matryoshka_vector_dimension=dimension,
+                            vector_dimension=dimension,
                             replace_field=reindex,
                         )
         return brain
@@ -671,9 +648,7 @@ class Resource:
         # Upload to binary storage
         # Vector indexing
         if self.disable_vectors is False:
-            await self.get_fields(force=True)
-            for field_vectors in message.field_vectors:
-                await self._apply_extracted_vectors(field_vectors)
+            await self._apply_extracted_vectors(message.field_vectors)
 
         # Only uploading to binary storage
         for field_large_metadata in message.field_large_metadata:
@@ -857,55 +832,69 @@ class Resource:
 
         add_field_classifications(self.basic, field_metadata)
 
-    async def _apply_extracted_vectors(self, field_vectors: ExtractedVectorsWrapper):
-        # Store vectors in the resource
-
-        if not self.has_field(field_vectors.field.field_type, field_vectors.field.field):
-            # skipping because field does not exist
-            logger.warning(f'Field "{field_vectors.field.field}" does not exist, skipping vectors')
-            return
-
-        field_obj = await self.get_field(
-            field_vectors.field.field,
-            field_vectors.field.field_type,
-            load=False,
-        )
-        vo = await field_obj.set_vectors(field_vectors)
-
-        # Prepare vectors to be indexed
+    async def _apply_extracted_vectors(
+        self,
+        fields_vectors: Sequence[ExtractedVectorsWrapper],
+    ):
+        await self.get_fields(force=True)
+        vectorsets = {
+            vectorset_id: vs
+            async for vectorset_id, vs in datamanagers.vectorsets.iter(self.txn, kbid=self.kb.kbid)
+        }
+
+        for field_vectors in fields_vectors:
+            # Bw/c with extracted vectors without vectorsets
+            if not field_vectors.vectorset_id:
+                assert (
+                    len(vectorsets) == 1
+                ), "Invalid broker message, can't ingest vectors from unknown vectorset to KB with multiple vectorsets"
+                vectorset = list(vectorsets.values())[0]
 
-        field_key = self.generate_field_id(field_vectors.field)
-        if vo is not None:
-            vectorset_id = field_vectors.vectorset_id or None
-            if vectorset_id is None:
-                dimension = await datamanagers.kb.get_matryoshka_vector_dimension(
-                    self.txn, kbid=self.kb.kbid
-                )
             else:
-                config = await datamanagers.vectorsets.get(
-                    self.txn, kbid=self.kb.kbid, vectorset_id=vectorset_id
-                )
-                if config is None:
+                if field_vectors.vectorset_id not in vectorsets:
                     logger.warning(
-                        f"Trying to apply a resource on vectorset '{vectorset_id}' that doesn't exist."
+                        "Dropping extracted vectors for unknown vectorset",
+                        extra={"kbid": self.kb.kbid, "vectorset": field_vectors.vectorset_id},
                     )
-                    return
-                dimension = config.vectorset_index_config.vector_dimension
-                if not dimension:
-                    raise ValueError(f"Vector dimension not set for vectorset '{vectorset_id}'")
+                    continue
+
+                vectorset = vectorsets[field_vectors.vectorset_id]
+
+            # Store vectors in the resource
+
+            if not self.has_field(field_vectors.field.field_type, field_vectors.field.field):
+                # skipping because field does not exist
+                logger.warning(f'Field "{field_vectors.field.field}" does not exist, skipping vectors')
+                return
+
+            field_obj = await self.get_field(
+                field_vectors.field.field,
+                field_vectors.field.field_type,
+                load=False,
+            )
+            vo = await field_obj.set_vectors(
+                field_vectors, vectorset.vectorset_id, vectorset.storage_key_kind
+            )
+            if vo is None:
+                raise AttributeError("Vector object not found on set_vectors")
+
+            # Prepare vectors to be indexed
+
+            field_key = self.generate_field_id(field_vectors.field)
+            dimension = vectorset.vectorset_index_config.vector_dimension
+            if not dimension:
+                raise ValueError(f"Vector dimension not set for vectorset '{vectorset.vectorset_id}'")
 
             apply_field_vectors_partial = partial(
                 self.indexer.apply_field_vectors,
                 field_key,
                 vo,
-                vectorset=vectorset_id,
+                vectorset=vectorset.vectorset_id,
                 replace_field=True,
-                matryoshka_vector_dimension=dimension,
+                vector_dimension=dimension,
             )
             loop = asyncio.get_running_loop()
             await loop.run_in_executor(_executor, apply_field_vectors_partial)
-        else:
-            raise AttributeError("VO not found on set")
 
     async def _apply_field_large_metadata(self, field_large_metadata: LargeComputedMetadataWrapper):
         field_obj = await self.get_field(
@@ -978,291 +967,6 @@ class Resource:
         self._indexer = None
         self.txn = None
 
-    async def iterate_sentences(
-        self, enabled_metadata: EnabledMetadata
-    ) -> AsyncIterator[TrainSentence]:  # pragma: no cover
-        fields = await self.get_fields(force=True)
-        metadata = TrainMetadata()
-        userdefinedparagraphclass: dict[str, ParagraphAnnotation] = {}
-        if enabled_metadata.labels:
-            if self.basic is None:
-                self.basic = await self.get_basic()
-            if self.basic is not None:
-                metadata.labels.resource.extend(self.basic.usermetadata.classifications)
-                for fieldmetadata in self.basic.fieldmetadata:
-                    field_id = self.generate_field_id(fieldmetadata.field)
-                    for annotationparagraph in fieldmetadata.paragraphs:
-                        userdefinedparagraphclass[annotationparagraph.key] = annotationparagraph
-
-        for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)
-            field_key = self.generate_field_id(fieldid)
-            fm = await field.get_field_metadata()
-            extracted_text = None
-            vo = None
-            text = None
-
-            if enabled_metadata.vector:
-                vo = await field.get_vectors()
-
-            extracted_text = await field.get_extracted_text()
-
-            if fm is None:
-                continue
-
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
-            for subfield_metadata, splitted_metadata in fm.split_metadata.items():
-                field_metadatas.append((subfield_metadata, splitted_metadata))
-
-            for subfield, field_metadata in field_metadatas:
-                if enabled_metadata.labels:
-                    metadata.labels.ClearField("field")
-                    metadata.labels.field.extend(field_metadata.classifications)
-
-                entities: dict[str, str] = {}
-                if enabled_metadata.entities:
-                    _update_entities_dict(entities, field_metadata)
-
-                precomputed_vectors = {}
-                if vo is not None:
-                    if subfield is not None:
-                        vectors = vo.split_vectors[subfield]
-                        base_vector_key = f"{self.uuid}/{field_key}/{subfield}"
-                    else:
-                        vectors = vo.vectors
-                        base_vector_key = f"{self.uuid}/{field_key}"
-                    for index, vector in enumerate(vectors.vectors):
-                        vector_key = f"{base_vector_key}/{index}/{vector.start}-{vector.end}"
-                        precomputed_vectors[vector_key] = vector.vector
-
-                if extracted_text is not None:
-                    if subfield is not None:
-                        text = extracted_text.split_text[subfield]
-                    else:
-                        text = extracted_text.text
-
-                for paragraph in field_metadata.paragraphs:
-                    if subfield is not None:
-                        paragraph_key = (
-                            f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
-                        )
-                    else:
-                        paragraph_key = f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
-
-                    if enabled_metadata.labels:
-                        metadata.labels.ClearField("field")
-                        metadata.labels.paragraph.extend(paragraph.classifications)
-                        if paragraph_key in userdefinedparagraphclass:
-                            metadata.labels.paragraph.extend(
-                                userdefinedparagraphclass[paragraph_key].classifications
-                            )
-
-                    for index, sentence in enumerate(paragraph.sentences):
-                        if subfield is not None:
-                            sentence_key = f"{self.uuid}/{field_key}/{subfield}/{index}/{sentence.start}-{sentence.end}"
-                        else:
-                            sentence_key = (
-                                f"{self.uuid}/{field_key}/{index}/{sentence.start}-{sentence.end}"
-                            )
-
-                        if vo is not None:
-                            metadata.ClearField("vector")
-                            vector_tmp = precomputed_vectors.get(sentence_key)
-                            if vector_tmp:
-                                metadata.vector.extend(vector_tmp)
-
-                        if extracted_text is not None and text is not None:
-                            metadata.text = text[sentence.start : sentence.end]
-
-                        metadata.ClearField("entities")
-                        metadata.ClearField("entity_positions")
-                        if enabled_metadata.entities and text is not None:
-                            local_text = text[sentence.start : sentence.end]
-                            add_entities_to_metadata(entities, local_text, metadata)
-
-                        pb_sentence = TrainSentence()
-                        pb_sentence.uuid = self.uuid
-                        pb_sentence.field.CopyFrom(fieldid)
-                        pb_sentence.paragraph = paragraph_key
-                        pb_sentence.sentence = sentence_key
-                        pb_sentence.metadata.CopyFrom(metadata)
-                        yield pb_sentence
-
-    async def iterate_paragraphs(
-        self, enabled_metadata: EnabledMetadata
-    ) -> AsyncIterator[TrainParagraph]:
-        fields = await self.get_fields(force=True)
-        metadata = TrainMetadata()
-        userdefinedparagraphclass: dict[str, ParagraphAnnotation] = {}
-        if enabled_metadata.labels:
-            if self.basic is None:
-                self.basic = await self.get_basic()
-            if self.basic is not None:
-                metadata.labels.resource.extend(self.basic.usermetadata.classifications)
-                for fieldmetadata in self.basic.fieldmetadata:
-                    field_id = self.generate_field_id(fieldmetadata.field)
-                    for annotationparagraph in fieldmetadata.paragraphs:
-                        userdefinedparagraphclass[annotationparagraph.key] = annotationparagraph
-
-        for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)
-            field_key = self.generate_field_id(fieldid)
-            fm = await field.get_field_metadata()
-            extracted_text = None
-            text = None
-
-            extracted_text = await field.get_extracted_text()
-
-            if fm is None:
-                continue
-
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
-            for subfield_metadata, splitted_metadata in fm.split_metadata.items():
-                field_metadatas.append((subfield_metadata, splitted_metadata))
-
-            for subfield, field_metadata in field_metadatas:
-                if enabled_metadata.labels:
-                    metadata.labels.ClearField("field")
-                    metadata.labels.field.extend(field_metadata.classifications)
-
-                entities: dict[str, str] = {}
-                if enabled_metadata.entities:
-                    _update_entities_dict(entities, field_metadata)
-
-                if extracted_text is not None:
-                    if subfield is not None:
-                        text = extracted_text.split_text[subfield]
-                    else:
-                        text = extracted_text.text
-
-                for paragraph in field_metadata.paragraphs:
-                    if subfield is not None:
-                        paragraph_key = (
-                            f"{self.uuid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
-                        )
-                    else:
-                        paragraph_key = f"{self.uuid}/{field_key}/{paragraph.start}-{paragraph.end}"
-
-                    if enabled_metadata.labels:
-                        metadata.labels.ClearField("paragraph")
-                        metadata.labels.paragraph.extend(paragraph.classifications)
-
-                        if extracted_text is not None and text is not None:
-                            metadata.text = text[paragraph.start : paragraph.end]
-
-                        metadata.ClearField("entities")
-                        metadata.ClearField("entity_positions")
-                        if enabled_metadata.entities and text is not None:
-                            local_text = text[paragraph.start : paragraph.end]
-                            add_entities_to_metadata(entities, local_text, metadata)
-
-                        if paragraph_key in userdefinedparagraphclass:
-                            metadata.labels.paragraph.extend(
-                                userdefinedparagraphclass[paragraph_key].classifications
-                            )
-
-                        pb_paragraph = TrainParagraph()
-                        pb_paragraph.uuid = self.uuid
-                        pb_paragraph.field.CopyFrom(fieldid)
-                        pb_paragraph.paragraph = paragraph_key
-                        pb_paragraph.metadata.CopyFrom(metadata)
-
-                        yield pb_paragraph
-
-    async def iterate_fields(self, enabled_metadata: EnabledMetadata) -> AsyncIterator[TrainField]:
-        fields = await self.get_fields(force=True)
-        metadata = TrainMetadata()
-        if enabled_metadata.labels:
-            if self.basic is None:
-                self.basic = await self.get_basic()
-            if self.basic is not None:
-                metadata.labels.resource.extend(self.basic.usermetadata.classifications)
-
-        for (type_id, field_id), field in fields.items():
-            fieldid = FieldID(field_type=type_id, field=field_id)
-            fm = await field.get_field_metadata()
-            extracted_text = None
-
-            if enabled_metadata.text:
-                extracted_text = await field.get_extracted_text()
-
-            if fm is None:
-                continue
-
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
-            for subfield_metadata, splitted_metadata in fm.split_metadata.items():
-                field_metadatas.append((subfield_metadata, splitted_metadata))
-
-            for subfield, splitted_metadata in field_metadatas:
-                if enabled_metadata.labels:
-                    metadata.labels.ClearField("field")
-                    metadata.labels.field.extend(splitted_metadata.classifications)
-
-                if extracted_text is not None:
-                    if subfield is not None:
-                        metadata.text = extracted_text.split_text[subfield]
-                    else:
-                        metadata.text = extracted_text.text
-
-                if enabled_metadata.entities:
-                    metadata.ClearField("entities")
-                    _update_entities_dict(metadata.entities, splitted_metadata)
-
-                pb_field = TrainField()
-                pb_field.uuid = self.uuid
-                pb_field.field.CopyFrom(fieldid)
-                pb_field.metadata.CopyFrom(metadata)
-                yield pb_field
-
-    async def generate_train_resource(self, enabled_metadata: EnabledMetadata) -> TrainResource:
-        fields = await self.get_fields(force=True)
-        metadata = TrainMetadata()
-        if enabled_metadata.labels:
-            if self.basic is None:
-                self.basic = await self.get_basic()
-            if self.basic is not None:
-                metadata.labels.resource.extend(self.basic.usermetadata.classifications)
-
-        metadata.labels.ClearField("field")
-        metadata.ClearField("entities")
-
-        for (_, _), field in fields.items():
-            extracted_text = None
-            fm = await field.get_field_metadata()
-
-            if enabled_metadata.text:
-                extracted_text = await field.get_extracted_text()
-
-            if extracted_text is not None:
-                metadata.text += extracted_text.text
-                for text in extracted_text.split_text.values():
-                    metadata.text += f" {text}"
-
-            if fm is None:
-                continue
-
-            field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
-            for subfield_metadata, splitted_metadata in fm.split_metadata.items():
-                field_metadatas.append((subfield_metadata, splitted_metadata))
-
-            for _, splitted_metadata in field_metadatas:
-                if enabled_metadata.labels:
-                    metadata.labels.field.extend(splitted_metadata.classifications)
-
-                if enabled_metadata.entities:
-                    _update_entities_dict(metadata.entities, splitted_metadata)
-
-        pb_resource = TrainResource()
-        pb_resource.uuid = self.uuid
-        if self.basic is not None:
-            pb_resource.title = self.basic.title
-            pb_resource.icon = self.basic.icon
-            pb_resource.slug = self.basic.slug
-            pb_resource.modified.CopyFrom(self.basic.modified)
-            pb_resource.created.CopyFrom(self.basic.created)
-        pb_resource.metadata.CopyFrom(metadata)
-        return pb_resource
-
 
 async def get_file_page_positions(field: File) -> FilePagePositions:
     positions: FilePagePositions = {}
@@ -1307,24 +1011,6 @@ def add_field_classifications(basic: PBBasic, fcmw: FieldComputedMetadataWrapper
     return True
 
 
-def add_entities_to_metadata(entities: dict[str, str], local_text: str, metadata: TrainMetadata) -> None:
-    for entity_key, entity_value in entities.items():
-        if entity_key not in local_text:
-            # Add the entity only if found in text
-            continue
-        metadata.entities[entity_key] = entity_value
-
-        # Add positions for the entity relative to the local text
-        poskey = f"{entity_value}/{entity_key}"
-        metadata.entity_positions[poskey].entity = entity_key
-        last_occurrence_end = 0
-        for _ in range(local_text.count(entity_key)):
-            start = local_text.index(entity_key, last_occurrence_end)
-            end = start + len(entity_key)
-            metadata.entity_positions[poskey].positions.append(TrainPosition(start=start, end=end))
-            last_occurrence_end = end
-
-
 def maybe_update_basic_summary(basic: PBBasic, summary_text: str) -> bool:
     if basic.summary or not summary_text:
         return False
@@ -1393,23 +1079,3 @@ def extract_field_metadata_languages(
     for _, splitted_metadata in field_metadata.metadata.split_metadata.items():
         languages.add(splitted_metadata.language)
     return list(languages)
-
-
-def _update_entities_dict(target_entites_dict: MutableMapping[str, str], field_metadata: FieldMetadata):
-    """
-    Update the entities dict with the entities from the field metadata.
-    Method created to ease the transition from legacy ner field to new entities field.
-    """
-    # Data Augmentation + Processor entities
-    # This will overwrite entities detected from more than one data augmentation task
-    # TODO: Change TrainMetadata proto to accept multiple entities with the same text
-    entity_map = {
-        entity.text: entity.label
-        for data_augmentation_task_id, entities_wrapper in field_metadata.entities.items()
-        for entity in entities_wrapper.entities
-    }
-    target_entites_dict.update(entity_map)
-
-    # Legacy processor entities
-    # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
-    target_entites_dict.update(field_metadata.ner)

nucliadb/ingest/serialize.py

@@ -17,10 +17,10 @@
 # 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 Optional
 
 import nucliadb_models as models
+from nucliadb.common import datamanagers
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.common.models_utils import from_proto
@@ -83,7 +83,18 @@ async def set_resource_field_extracted_data(
             field_data.large_metadata = from_proto.large_computed_metadata(data_lcm)
 
     if ExtractedDataTypeName.VECTOR in wanted_extracted_data:
-        data_vec = await field.get_vectors()
+        # 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)
 

nucliadb/ingest/service/writer.py

@@ -485,6 +485,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
                 vector_dimension=request.vector_dimension,
             ),
             matryoshka_dimensions=request.matryoshka_dimensions,
+            storage_key_kind=VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX,
         )
         response = NewVectorSetResponse()
         try:
@@ -513,6 +514,9 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
                 kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
                 await kbobj.delete_vectorset(request.vectorset_id)
                 await txn.commit()
+        except VectorSetConflict as exc:
+            response.status = DelVectorSetResponse.Status.ERROR
+            response.details = str(exc)
         except Exception as exc:
             errors.capture_exception(exc)
             logger.error("Error in ingest gRPC while deleting a vectorset", exc_info=True)

nucliadb/purge/__init__.py

@@ -26,6 +26,7 @@ from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
 from nucliadb.ingest import SERVICE_NAME, logger
+from nucliadb.ingest.fields.base import Field
 from nucliadb.ingest.orm.knowledgebox import (
     KB_TO_DELETE,
     KB_TO_DELETE_BASE,
@@ -35,6 +36,7 @@ from nucliadb.ingest.orm.knowledgebox import (
     RESOURCE_TO_DELETE_STORAGE_BASE,
     KnowledgeBox,
 )
+from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig, VectorSetPurge
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
 from nucliadb_utils.storages.storage import Storage
@@ -201,8 +203,8 @@ async def purge_kb_vectorsets(driver: Driver, storage: Storage):
     """
     logger.info("START PURGING KB VECTORSETS")
 
-    purged = []
-    async for key in _iter_keys(driver, KB_VECTORSET_TO_DELETE_BASE):
+    vectorsets_to_delete = [key async for key in _iter_keys(driver, KB_VECTORSET_TO_DELETE_BASE)]
+    for key in vectorsets_to_delete:
         logger.info(f"Purging vectorsets {key}")
         try:
             _base, kbid, vectorset = key.lstrip("/").split("/")
@@ -211,13 +213,38 @@ async def purge_kb_vectorsets(driver: Driver, storage: Storage):
             continue
 
         try:
+            async with driver.transaction(read_only=True) as txn:
+                value = await txn.get(key)
+                assert value is not None, "Key must exist or we wouldn't had fetch it iterating keys"
+                purge_payload = VectorSetPurge()
+                purge_payload.ParseFromString(value)
+
+            fields: list[Field] = []
             async with driver.transaction(read_only=True) as txn:
                 kb = KnowledgeBox(txn, storage, kbid)
                 async for resource in kb.iterate_resources():
-                    fields = await resource.get_fields(force=True)
+                    fields.extend((await resource.get_fields(force=True)).values())
+
             # we don't need the maindb transaction anymore to remove vectors from storage
-            for field in fields.values():
-                await field.delete_vectors(vectorset)
+            for field in fields:
+                if purge_payload.storage_key_kind == VectorSetConfig.StorageKeyKind.UNSET:
+                    # Bw/c for purge before adding purge payload. We assume
+                    # there's only 2 kinds of KBs: with one or with more than
+                    # one vectorset. KBs with one vectorset are not allowed to
+                    # delete their vectorset, so we wouldn't be here. It has to
+                    # be a KB with multiple, so the storage key kind has to be
+                    # this:
+                    await field.delete_vectors(
+                        vectorset, VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX
+                    )
+                else:
+                    await field.delete_vectors(vectorset, purge_payload.storage_key_kind)
+
+            # Finally, delete the key
+            async with driver.transaction() as txn:
+                await txn.delete(key)
+                await txn.commit()
+
         except Exception as exc:
             errors.capture_exception(exc)
             logger.error(
@@ -227,13 +254,6 @@ async def purge_kb_vectorsets(driver: Driver, storage: Storage):
             )
             continue
 
-        purged.append(key)
-
-    async with driver.transaction() as txn:
-        for key in purged:
-            await txn.delete(key)
-        await txn.commit()
-
     logger.info("FINISH PURGING KB VECTORSETS")