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

nucliadb/ingest/processing.py

@@ -261,6 +261,7 @@ class ProcessingEngine:
             "content_type": file.file.content_type,
             "password": file.password,
             "language": file.language,
+            "extract_strategy": file.extract_strategy,
         }
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
@@ -278,6 +279,8 @@ class ProcessingEngine:
         headers["X-LANGUAGE"] = file.language
         headers["X-FILENAME"] = base64.b64encode(file.file.filename.encode()).decode()  # type: ignore
         headers["X-MD5"] = file.file.md5
+        if file.extract_strategy is not None:
+            headers["X-EXTRACT-STRATEGY"] = file.extract_strategy
         headers["CONTENT_TYPE"] = file.file.content_type
         headers["CONTENT-LENGTH"] = str(len(file.file.payload))  # type: ignore
         headers["X-STF-NUAKEY"] = f"Bearer {self.nuclia_service_account}"
@@ -317,6 +320,7 @@ class ProcessingEngine:
             "content_type": file_field.file.content_type,
             "language": file_field.language,
             "password": file_field.password,
+            "extract_strategy": file_field.extract_strategy,
         }
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
@@ -341,6 +345,8 @@ class ProcessingEngine:
             headers["CONTENT-TYPE"] = file.file.content_type
             if file.file.size:
                 headers["CONTENT-LENGTH"] = str(file.file.size)
+            if file.extract_strategy != "":
+                headers["X-EXTRACT-STRATEGY"] = file.extract_strategy
             headers["X-STF-NUAKEY"] = f"Bearer {self.nuclia_service_account}"
 
             iterator = storage.downloadbytescf_iterator(file.file)
@@ -406,13 +412,13 @@ class ProcessingEngine:
                 # Upload the payload
                 item.partition = partition
                 resp = await self.session.post(
-                    url=self.nuclia_internal_push, data=item.json(), headers=headers
+                    url=self.nuclia_internal_push, data=item.model_dump_json(), headers=headers
                 )
             else:
                 headers.update({"X-STF-NUAKEY": f"Bearer {self.nuclia_service_account}"})
                 # Upload the payload
                 resp = await self.session.post(
-                    url=self.nuclia_external_push_v2, data=item.json(), headers=headers
+                    url=self.nuclia_external_push_v2, data=item.model_dump_json(), headers=headers
                 )
             if resp.status == 200:
                 data = await resp.json()

nucliadb/ingest/serialize.py

@@ -18,18 +18,20 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
+from typing import Optional, Union
 
 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
 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.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.resource import Resource as ORMResource
-from nucliadb_models.common import FIELD_TYPES_MAP, FieldTypeName
+from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import (
     ConversationFieldData,
     ConversationFieldExtractedData,
@@ -49,7 +51,9 @@ from nucliadb_models.resource import (
 )
 from nucliadb_models.search import ResourceProperties
 from nucliadb_models.security import ResourceSecurity
-from nucliadb_utils.utilities import get_storage
+from nucliadb_protos.writer_pb2 import FieldStatus
+from nucliadb_utils import const
+from nucliadb_utils.utilities import get_storage, has_feature
 
 
 async def set_resource_field_extracted_data(
@@ -64,7 +68,7 @@ async def set_resource_field_extracted_data(
     if ExtractedDataTypeName.TEXT in wanted_extracted_data:
         data_et = await field.get_extracted_text()
         if data_et is not None:
-            field_data.text = models.ExtractedText.from_message(data_et)
+            field_data.text = from_proto.extracted_text(data_et)
 
     metadata_wanted = ExtractedDataTypeName.METADATA in wanted_extracted_data
     shortened_metadata_wanted = ExtractedDataTypeName.SHORTENED_METADATA in wanted_extracted_data
@@ -72,24 +76,35 @@ async def set_resource_field_extracted_data(
         data_fcm = await field.get_field_metadata()
 
         if data_fcm is not None:
-            field_data.metadata = models.FieldComputedMetadata.from_message(
+            field_data.metadata = from_proto.field_computed_metadata(
                 data_fcm, shortened=shortened_metadata_wanted and not metadata_wanted
             )
 
     if ExtractedDataTypeName.LARGE_METADATA in wanted_extracted_data:
         data_lcm = await field.get_large_field_metadata()
         if data_lcm is not None:
-            field_data.large_metadata = models.LargeComputedMetadata.from_message(data_lcm)
+            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 = models.VectorObject.from_message(data_vec)
+            field_data.vectors = from_proto.vector_object(data_vec)
 
     if ExtractedDataTypeName.QA in wanted_extracted_data:
         qa = await field.get_question_answers()
         if qa is not None:
-            field_data.question_answers = models.FieldQuestionAnswers.from_message(qa)
+            field_data.question_answers = from_proto.field_question_answers(qa)
 
     if (
         isinstance(field, File)
@@ -98,7 +113,7 @@ async def set_resource_field_extracted_data(
     ):
         data_fed = await field.get_file_extracted_data()
         if data_fed is not None:
-            field_data.file = models.FileExtractedData.from_message(data_fed)
+            field_data.file = from_proto.file_extracted_data(data_fed)
 
     if (
         isinstance(field, Link)
@@ -107,7 +122,7 @@ async def set_resource_field_extracted_data(
     ):
         data_led = await field.get_link_extracted_data()
         if data_led is not None:
-            field_data.link = models.LinkExtractedData.from_message(data_led)
+            field_data.link = from_proto.link_extracted_data(data_led)
 
 
 async def serialize(
@@ -133,6 +148,40 @@ async def serialize(
         )
 
 
+async def serialize_field_errors(
+    field: Field,
+    serialized: Union[
+        TextFieldData, FileFieldData, LinkFieldData, ConversationFieldData, GenericFieldData
+    ],
+):
+    if has_feature(const.Features.FIELD_STATUS):
+        status = await field.get_status()
+        if status is None:
+            status = FieldStatus()
+        serialized.status = status.Status.Name(status.status)
+        if status.errors:
+            serialized.errors = []
+            for error in status.errors:
+                serialized.errors.append(
+                    Error(
+                        body=error.source_error.error,
+                        code=error.source_error.code,
+                        code_str=error.source_error.ErrorCode.Name(error.source_error.code),
+                        created=error.created.ToDatetime(),
+                    )
+                )
+            serialized.error = serialized.errors[-1]
+    else:
+        field_error = await field.get_error()
+        if field_error is not None:
+            serialized.error = Error(
+                body=field_error.error,
+                code=field_error.code,
+                code_str=field_error.ErrorCode.Name(field_error.code),
+                created=None,
+            )
+
+
 async def managed_serialize(
     txn: Transaction,
     kbid: str,
@@ -174,14 +223,12 @@ async def managed_serialize(
                 else None
             )
 
-            resource.metadata = models.Metadata.from_message(orm_resource.basic.metadata)
-            resource.usermetadata = models.UserMetadata.from_message(orm_resource.basic.usermetadata)
+            resource.metadata = from_proto.metadata(orm_resource.basic.metadata)
+            resource.usermetadata = from_proto.user_metadata(orm_resource.basic.usermetadata)
             resource.fieldmetadata = [
-                models.UserFieldMetadata.from_message(fm) for fm in orm_resource.basic.fieldmetadata
+                from_proto.user_field_metadata(fm) for fm in orm_resource.basic.fieldmetadata
             ]
-            resource.computedmetadata = models.ComputedMetadata.from_message(
-                orm_resource.basic.computedmetadata
-            )
+            resource.computedmetadata = from_proto.computed_metadata(orm_resource.basic.computedmetadata)
 
             resource.last_seqid = orm_resource.basic.last_seqid
 
@@ -195,18 +242,18 @@ async def managed_serialize(
         await orm_resource.get_relations()
         if orm_resource.relations is not None:
             resource.relations = [
-                models.Relation.from_message(relation) for relation in orm_resource.relations.relations
+                from_proto.relation(relation) for relation in orm_resource.relations.relations
             ]
 
     if ResourceProperties.ORIGIN in show:
         await orm_resource.get_origin()
         if orm_resource.origin is not None:
-            resource.origin = models.Origin.from_message(orm_resource.origin)
+            resource.origin = from_proto.origin(orm_resource.origin)
 
     if ResourceProperties.EXTRA in show:
         await orm_resource.get_extra()
         if orm_resource.extra is not None:
-            resource.extra = models.Extra.from_message(orm_resource.extra)
+            resource.extra = from_proto.extra(orm_resource.extra)
 
     include_errors = ResourceProperties.ERRORS in show
 
@@ -221,7 +268,7 @@ async def managed_serialize(
         await orm_resource.get_fields()
         resource.data = ResourceData()
         for (field_type, _), field in orm_resource.fields.items():
-            field_type_name = FIELD_TYPES_MAP[field_type]
+            field_type_name = from_proto.field_type_name(field_type)
             if field_type_name not in field_type_filter:
                 continue
 
@@ -236,14 +283,10 @@ async def managed_serialize(
                 if field.id not in resource.data.texts:
                     resource.data.texts[field.id] = TextFieldData()
                 if include_value:
-                    serialized_value = (
-                        models.FieldText.from_message(value) if value is not None else None
-                    )
+                    serialized_value = from_proto.field_text(value) if value is not None else None
                     resource.data.texts[field.id].value = serialized_value
                 if include_errors:
-                    error = await field.get_error()
-                    if error is not None:
-                        resource.data.texts[field.id].error = Error(body=error.error, code=error.code)
+                    await serialize_field_errors(field, resource.data.texts[field.id])
                 if include_extracted_data:
                     resource.data.texts[field.id].extracted = TextFieldExtractedData()
                     await set_resource_field_extracted_data(
@@ -259,14 +302,12 @@ async def managed_serialize(
                     resource.data.files[field.id] = FileFieldData()
                 if include_value:
                     if value is not None:
-                        resource.data.files[field.id].value = models.FieldFile.from_message(value)
+                        resource.data.files[field.id].value = from_proto.field_file(value)
                     else:
                         resource.data.files[field.id].value = None
 
                 if include_errors:
-                    error = await field.get_error()
-                    if error is not None:
-                        resource.data.files[field.id].error = Error(body=error.error, code=error.code)
+                    await serialize_field_errors(field, resource.data.files[field.id])
 
                 if include_extracted_data:
                     resource.data.files[field.id].extracted = FileFieldExtractedData()
@@ -282,12 +323,10 @@ async def managed_serialize(
                 if field.id not in resource.data.links:
                     resource.data.links[field.id] = LinkFieldData()
                 if include_value and value is not None:
-                    resource.data.links[field.id].value = models.FieldLink.from_message(value)
+                    resource.data.links[field.id].value = from_proto.field_link(value)
 
                 if include_errors:
-                    error = await field.get_error()
-                    if error is not None:
-                        resource.data.links[field.id].error = Error(body=error.error, code=error.code)
+                    await serialize_field_errors(field, resource.data.links[field.id])
 
                 if include_extracted_data:
                     resource.data.links[field.id].extracted = LinkFieldExtractedData()
@@ -303,16 +342,10 @@ async def managed_serialize(
                 if field.id not in resource.data.conversations:
                     resource.data.conversations[field.id] = ConversationFieldData()
                 if include_errors:
-                    error = await field.get_error()
-                    if error is not None:
-                        resource.data.conversations[field.id].error = Error(
-                            body=error.error, code=error.code
-                        )
+                    await serialize_field_errors(field, resource.data.conversations[field.id])
                 if include_value and isinstance(field, Conversation):
                     value = await field.get_metadata()
-                    resource.data.conversations[field.id].value = models.FieldConversation.from_message(
-                        value
-                    )
+                    resource.data.conversations[field.id].value = from_proto.field_conversation(value)
                 if include_extracted_data:
                     resource.data.conversations[field.id].extracted = ConversationFieldExtractedData()
                     await set_resource_field_extracted_data(
@@ -329,9 +362,7 @@ async def managed_serialize(
                 if include_value:
                     resource.data.generics[field.id].value = value
                 if include_errors:
-                    error = await field.get_error()
-                    if error is not None:
-                        resource.data.generics[field.id].error = Error(body=error.error, code=error.code)
+                    await serialize_field_errors(field, resource.data.generics[field.id])
                 if include_extracted_data:
                     resource.data.generics[field.id].extracted = TextFieldExtractedData(
                         text=models.ExtractedText(text=resource.data.generics[field.id].value)

nucliadb/ingest/service/writer.py

@@ -31,12 +31,12 @@ from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.ingest import SERVICE_NAME, logger
 from nucliadb.ingest.orm.broker_message import generate_broker_message
 from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict, VectorSetConflict
+from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
 from nucliadb.ingest.orm.processor import Processor, sequence_manager
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
 from nucliadb.ingest.settings import settings
-from nucliadb_protos import nodewriter_pb2, writer_pb2, writer_pb2_grpc
+from nucliadb_protos import writer_pb2, writer_pb2_grpc
 from nucliadb_protos.knowledgebox_pb2 import (
     DeleteKnowledgeBoxResponse,
     KnowledgeBoxID,
@@ -44,13 +44,10 @@ from nucliadb_protos.knowledgebox_pb2 import (
     KnowledgeBoxUpdate,
     SemanticModelMetadata,
     UpdateKnowledgeBoxResponse,
-    VectorSetConfig,
 )
 from nucliadb_protos.writer_pb2 import (
     BrokerMessage,
     DelEntitiesRequest,
-    DelVectorSetRequest,
-    DelVectorSetResponse,
     GetEntitiesGroupRequest,
     GetEntitiesGroupResponse,
     GetEntitiesRequest,
@@ -63,8 +60,6 @@ from nucliadb_protos.writer_pb2 import (
     ListMembersResponse,
     NewEntitiesGroupRequest,
     NewEntitiesGroupResponse,
-    NewVectorSetRequest,
-    NewVectorSetResponse,
     OpStatusWriter,
     SetEntitiesRequest,
     UpdateEntitiesGroupRequest,
@@ -472,52 +467,3 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             errors.capture_exception(e)
             logger.error("Error in ingest gRPC servicer", exc_info=True)
             raise
-
-    async def NewVectorSet(  # type: ignore
-        self, request: NewVectorSetRequest, context=None
-    ) -> NewVectorSetResponse:
-        config = VectorSetConfig(
-            vectorset_id=request.vectorset_id,
-            vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
-                similarity=request.similarity,
-                normalize_vectors=request.normalize_vectors,
-                vector_type=request.vector_type,
-                vector_dimension=request.vector_dimension,
-            ),
-            matryoshka_dimensions=request.matryoshka_dimensions,
-        )
-        response = NewVectorSetResponse()
-        try:
-            async with self.driver.transaction() as txn:
-                kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-                await kbobj.create_vectorset(config)
-                await txn.commit()
-        except VectorSetConflict as exc:
-            response.status = NewVectorSetResponse.Status.ERROR
-            response.details = str(exc)
-        except Exception as exc:
-            errors.capture_exception(exc)
-            logger.error("Error in ingest gRPC while creating a vectorset", exc_info=True)
-            response.status = NewVectorSetResponse.Status.ERROR
-            response.details = str(exc)
-        else:
-            response.status = NewVectorSetResponse.Status.OK
-        return response
-
-    async def DelVectorSet(  # type: ignore
-        self, request: DelVectorSetRequest, context=None
-    ) -> DelVectorSetResponse:
-        response = DelVectorSetResponse()
-        try:
-            async with self.driver.transaction() as txn:
-                kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-                await kbobj.delete_vectorset(request.vectorset_id)
-                await txn.commit()
-        except Exception as exc:
-            errors.capture_exception(exc)
-            logger.error("Error in ingest gRPC while deleting a vectorset", exc_info=True)
-            response.status = DelVectorSetResponse.Status.ERROR
-            response.details = str(exc)
-        else:
-            response.status = DelVectorSetResponse.Status.OK
-        return response

nucliadb/ingest/settings.py

@@ -77,7 +77,7 @@ class Settings(DriverSettings):
     total_replicas: int = 1  # number of ingest processor replicas in the cluster
     nuclia_partitions: int = 50
 
-    max_receive_message_length: int = 4
+    max_receive_message_length: int = 500  # In MB
 
     # Search query timeouts
     relation_search_timeout: float = 10.0
@@ -85,8 +85,5 @@ class Settings(DriverSettings):
 
     max_concurrent_ingest_processing: int = 5
 
-    # Ingest processor settings
-    ingest_delete_resource_storage_max_parallel: int = 20
-
 
 settings = Settings()

nucliadb/learning_proxy.py

@@ -155,10 +155,9 @@ class LearningConfiguration(BaseModel):
 
 
 class ProxiedLearningConfigError(Exception):
-    def __init__(self, status_code: int, content: bytes, content_type: str):
+    def __init__(self, status_code: int, content: Union[str, dict[str, Any]]):
         self.status_code = status_code
         self.content = content
-        self.content_type = content_type
 
 
 def raise_for_status(response: httpx.Response) -> None:
@@ -166,10 +165,13 @@ def raise_for_status(response: httpx.Response) -> None:
         response.raise_for_status()
     except httpx.HTTPStatusError as err:
         content_type = err.response.headers.get("Content-Type", "application/json")
+        if content_type == "application/json":
+            content = err.response.json()
+        else:
+            content = err.response.text
         raise ProxiedLearningConfigError(
             status_code=err.response.status_code,
-            content=err.response.content,
-            content_type=content_type,
+            content=content,
         )
 
 

nucliadb/purge/__init__.py

@@ -25,15 +25,19 @@ from nucliadb.common.cluster.exceptions import NodeError, ShardNotFound
 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.common.nidx import start_nidx_utility, stop_nidx_utility
 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,
     KB_TO_DELETE_STORAGE_BASE,
     KB_VECTORSET_TO_DELETE,
     KB_VECTORSET_TO_DELETE_BASE,
+    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
@@ -131,6 +135,67 @@ async def purge_kb_storage(driver: Driver, storage: Storage):
     logger.info("FINISH PURGING KB STORAGE")
 
 
+async def purge_deleted_resource_storage(driver: Driver, storage: Storage) -> None:
+    """
+    Remove from storage all resources marked as deleted.
+
+    Returns the number of resources purged.
+    """
+    logger.info("Starting purge of deleted resource storage")
+    to_purge = await _count_resources_storage_to_purge(driver)
+    logger.info(f"Found {to_purge} resources to purge")
+    while True:
+        try:
+            purged = await _purge_resources_storage_batch(driver, storage, batch_size=100)
+            if not purged:
+                logger.info("No more resources to purge found")
+                return
+            logger.info(f"Purged {purged} resources")
+
+        except asyncio.CancelledError:
+            logger.info("Purge of deleted resource storage was cancelled")
+            return
+
+
+async def _count_resources_storage_to_purge(driver: Driver) -> int:
+    """
+    Count the number of resources marked as deleted in storage.
+    """
+    async with driver.transaction(read_only=True) as txn:
+        return await txn.count(match=RESOURCE_TO_DELETE_STORAGE_BASE)
+
+
+async def _purge_resources_storage_batch(driver: Driver, storage: Storage, batch_size: int = 100) -> int:
+    """
+    Remove from storage a batch of resources marked as deleted. Returns the
+    number of resources purged.
+    """
+    # Get the keys of the resources to delete in batches of 100
+    to_delete_batch = []
+    async with driver.transaction(read_only=True) as txn:
+        async for key in txn.keys(match=RESOURCE_TO_DELETE_STORAGE_BASE, count=batch_size):
+            to_delete_batch.append(key)
+
+    if not to_delete_batch:
+        return 0
+
+    # Delete the resources blobs from storage
+    logger.info(f"Purging {len(to_delete_batch)} deleted resources")
+    tasks = []
+    for key in to_delete_batch:
+        kbid, resource_id = key.split("/")[-2:]
+        tasks.append(asyncio.create_task(storage.delete_resource(kbid, resource_id)))
+    await asyncio.gather(*tasks)
+
+    # Delete the schedule-to-delete keys
+    async with driver.transaction() as txn:
+        for key in to_delete_batch:
+            await txn.delete(key)
+        await txn.commit()
+
+    return len(to_delete_batch)
+
+
 async def purge_kb_vectorsets(driver: Driver, storage: Storage):
     """Vectors for a vectorset are stored in a key inside each resource. Iterate
     through all resources of the KB and remove any storage object containing
@@ -139,8 +204,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("/")
@@ -149,13 +214,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(
@@ -165,13 +255,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")
 
 
@@ -180,22 +263,29 @@ async def main():
     This script will purge all knowledge boxes marked to be deleted in maindb.
     """
     await setup_cluster()
+    await start_nidx_utility()
     driver = await setup_driver()
     storage = await get_storage(
         gcs_scopes=["https://www.googleapis.com/auth/devstorage.full_control"],
         service_name=SERVICE_NAME,
     )
     try:
+        purge_resources_storage_task = asyncio.create_task(
+            purge_deleted_resource_storage(driver, storage)
+        )
         await purge_kb(driver)
         await purge_kb_storage(driver, storage)
         await purge_kb_vectorsets(driver, storage)
+        await purge_resources_storage_task
     except Exception as ex:  # pragma: no cover
         logger.exception("Unhandled exception on purge command")
         errors.capture_exception(ex)
     finally:
         try:
+            purge_resources_storage_task.cancel()
             await storage.finalize()
             await teardown_driver()
+            await stop_nidx_utility()
             await teardown_cluster()
         except Exception:  # pragma: no cover
             logger.exception("Error tearing down utilities on purge command")

nucliadb/purge/orphan_shards.py

@@ -33,6 +33,7 @@ from nucliadb.common.cluster.manager import KBShardManager
 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.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.ingest import logger
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
@@ -135,10 +136,9 @@ async def _get_stored_shards(driver: Driver) -> dict[str, ShardLocation]:
                 continue
             else:
                 for shard_object_pb in kb_shards:
-                    for shard_replica_pb in shard_object_pb.replicas:
-                        shard_replica_id = shard_replica_pb.shard.id
-                        node_id = shard_replica_pb.node
-                        stored_shards[shard_replica_id] = ShardLocation(kbid=kbid, node_id=node_id)
+                    stored_shards[shard_object_pb.nidx_shard_id] = ShardLocation(
+                        kbid=kbid, node_id="nidx"
+                    )
     return stored_shards
 
 
@@ -241,6 +241,7 @@ async def main():
     """
     args = parse_arguments()
 
+    await start_nidx_utility()
     await setup_cluster()
     driver = await setup_driver()
 
@@ -253,6 +254,7 @@ async def main():
     finally:
         await teardown_driver()
         await teardown_cluster()
+        await stop_nidx_utility()
 
 
 def run() -> int:  # pragma: no cover