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

nucliadb/common/datamanagers/vectorsets.py

@@ -58,6 +58,11 @@ async def iter(
         yield config.vectorset_id, config
 
 
+async def count(txn: Transaction, *, kbid: str) -> int:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    return len(kb_vectorsets.vectorsets)
+
+
 async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
     """Create or update a vectorset configuration"""
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
@@ -73,16 +78,20 @@ async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSet
     await txn.set(key, kb_vectorsets.SerializeToString())
 
 
-async def delete(txn: Transaction, *, kbid: str, vectorset_id: str):
+async def delete(
+    txn: Transaction, *, kbid: str, vectorset_id: str
+) -> Optional[knowledgebox_pb2.VectorSetConfig]:
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
     index = _find_vectorset(kb_vectorsets, vectorset_id)
     if index is None:
         # already deleted
-        return
+        return None
 
+    deleted = kb_vectorsets.vectorsets[index]
     del kb_vectorsets.vectorsets[index]
     key = KB_VECTORSETS.format(kbid=kbid)
     await txn.set(key, kb_vectorsets.SerializeToString())
+    return deleted
 
 
 # XXX At some point in the vectorset epic, we should make this key mandatory and

nucliadb/ingest/fields/base.py

@@ -21,12 +21,13 @@ from __future__ import annotations
 
 import enum
 from datetime import datetime
-from typing import Any, Generic, Optional, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, Optional, Type, TypeVar
 
 from google.protobuf.message import DecodeError, Message
 
 from nucliadb.common import datamanagers
 from nucliadb.ingest.fields.exceptions import InvalidFieldClass, InvalidPBClass
+from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import (
     CloudFile,
     ExtractedTextWrapper,
@@ -44,9 +45,15 @@ from nucliadb_protos.utils_pb2 import ExtractedText, VectorObject
 from nucliadb_protos.writer_pb2 import Error, FieldStatus
 from nucliadb_utils.storages.storage import Storage, StorageField
 
+if TYPE_CHECKING:  # pragma: no cover
+    from nucliadb.ingest.orm.resource import Resource
+
+
 SUBFIELDFIELDS = ("c",)
 
 
+# NOTE extracted vectors key is no longer a static key, it is stored in each
+# vectorset
 class FieldTypes(str, enum.Enum):
     FIELD_TEXT = "extracted_text"
     FIELD_VECTORS = "extracted_vectors"
@@ -73,7 +80,7 @@ class Field(Generic[PbType]):
     def __init__(
         self,
         id: str,
-        resource: Any,
+        resource: Resource,
         pb: Optional[Any] = None,
         value: Optional[Any] = None,
     ):
@@ -88,7 +95,7 @@ class Field(Generic[PbType]):
         self.question_answers = None
 
         self.id: str = id
-        self.resource: Any = resource
+        self.resource = resource
 
         if value is not None:
             newpb = self.pbklass()
@@ -119,11 +126,20 @@ class Field(Generic[PbType]):
     def get_storage_field(self, field_type: FieldTypes) -> StorageField:
         return self.storage.file_extracted(self.kbid, self.uuid, self.type, self.id, field_type.value)
 
-    def _get_extracted_vectors_storage_field(self, vectorset: Optional[str] = None) -> StorageField:
-        if vectorset:
+    def _get_extracted_vectors_storage_field(
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> StorageField:
+        if storage_key_kind == VectorSetConfig.StorageKeyKind.LEGACY:
+            key = FieldTypes.FIELD_VECTORS.value
+        elif storage_key_kind == VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX:
             key = FieldTypes.FIELD_VECTORSET.value.format(vectorset=vectorset)
         else:
-            key = FieldTypes.FIELD_VECTORS.value
+            raise ValueError(
+                f"Can't do anything with UNSET or unknown vectorset storage key kind: {storage_key_kind}"
+            )
+
         return self.storage.file_extracted(self.kbid, self.uuid, self.type, self.id, key)
 
     async def db_get_value(self) -> Optional[PbType]:
@@ -163,7 +179,8 @@ class Field(Generic[PbType]):
             field_id=self.id,
         )
         await self.delete_extracted_text()
-        await self.delete_vectors()
+        async for vectorset_id, vs in datamanagers.vectorsets.iter(self.resource.txn, kbid=self.kbid):
+            await self.delete_vectors(vectorset_id, vs.storage_key_kind)
         await self.delete_metadata()
         await self.delete_question_answers()
 
@@ -181,9 +198,13 @@ class Field(Generic[PbType]):
         except KeyError:
             pass
 
-    async def delete_vectors(self, vectorset: Optional[str] = None) -> None:
+    async def delete_vectors(
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> None:
         # Try delete vectors
-        sf = self._get_extracted_vectors_storage_field(vectorset)
+        sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
         try:
             await self.storage.delete_upload(sf.key, sf.bucket)
         except KeyError:
@@ -328,12 +349,17 @@ class Field(Generic[PbType]):
                 self.extracted_text = payload
         return self.extracted_text
 
-    async def set_vectors(self, payload: ExtractedVectorsWrapper) -> Optional[VectorObject]:
-        vectorset = payload.vectorset_id or None
+    async def set_vectors(
+        self,
+        payload: ExtractedVectorsWrapper,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+    ) -> Optional[VectorObject]:
         if self.type in SUBFIELDFIELDS:
             try:
                 actual_payload: Optional[VectorObject] = await self.get_vectors(
                     vectorset=vectorset,
+                    storage_key_kind=storage_key_kind,
                     force=True,
                 )
             except KeyError:
@@ -341,7 +367,7 @@ class Field(Generic[PbType]):
         else:
             actual_payload = None
 
-        sf = self._get_extracted_vectors_storage_field(vectorset)
+        sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
         vo: Optional[VectorObject] = None
         if actual_payload is None:
             # Its first extracted text
@@ -373,14 +399,13 @@ class Field(Generic[PbType]):
         return vo
 
     async def get_vectors(
-        self, vectorset: Optional[str] = None, force: bool = False
+        self,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
+        force: bool = False,
     ) -> Optional[VectorObject]:
-        # compat with vectorsets coming from protobuffers where no value is
-        # empty string instead of None. This shouldn't be handled here but we
-        # have to make sure it gets the correct vectorset
-        vectorset = vectorset or None
         if self.extracted_vectors.get(vectorset, None) is None or force:
-            sf = self._get_extracted_vectors_storage_field(vectorset)
+            sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
             payload = await self.storage.download_pb(sf, VectorObject)
             if payload is not None:
                 self.extracted_vectors[vectorset] = payload

nucliadb/ingest/orm/brain.py

@@ -246,9 +246,10 @@ class ResourceBrain:
         field_id: str,
         vo: utils_pb2.VectorObject,
         *,
-        vectorset: Optional[str] = None,
+        vectorset: str,
         replace_field: bool = False,
-        matryoshka_vector_dimension: Optional[int] = None,
+        # cut to specific dimension if specified
+        vector_dimension: Optional[int] = None,
     ):
         fid = ids.FieldId.from_string(f"{self.rid}/{field_id}")
         for subfield, vectors in vo.split_vectors.items():
@@ -277,7 +278,7 @@ class ResourceBrain:
                     sentence_key,
                     vector,
                     vectorset=vectorset,
-                    matryoshka_vector_dimension=matryoshka_vector_dimension,
+                    vector_dimension=vector_dimension,
                 )
 
         _field_id = ids.FieldId(
@@ -303,16 +304,12 @@ class ResourceBrain:
                 sentence_key,
                 vector,
                 vectorset=vectorset,
-                matryoshka_vector_dimension=matryoshka_vector_dimension,
+                vector_dimension=vector_dimension,
             )
 
         if replace_field:
             full_field_id = ids.FieldId(rid=self.rid, type=fid.type, key=fid.key).full()
-            if vectorset is None:
-                # DEPRECATED
-                self.brain.sentences_to_delete.append(full_field_id)
-            else:
-                self.brain.vector_prefixes_to_delete[vectorset].items.append(full_field_id)
+            self.brain.vector_prefixes_to_delete[vectorset].items.append(full_field_id)
 
     def _apply_field_vector(
         self,
@@ -321,22 +318,15 @@ class ResourceBrain:
         sentence_key: ids.VectorId,
         vector: utils_pb2.Vector,
         *,
-        vectorset: Optional[str],
-        matryoshka_vector_dimension: Optional[int] = None,
+        vectorset: str,
+        # cut vectors if a specific dimension is specified
+        vector_dimension: Optional[int] = None,
     ):
         paragraph_pb = self.brain.paragraphs[field_id].paragraphs[paragraph_key.full()]
-        if vectorset:
-            sentence_pb = paragraph_pb.vectorsets_sentences[vectorset].sentences[sentence_key.full()]
-        else:
-            sentence_pb = paragraph_pb.sentences[sentence_key.full()]
+        sentence_pb = paragraph_pb.vectorsets_sentences[vectorset].sentences[sentence_key.full()]
 
         sentence_pb.ClearField("vector")  # clear first to prevent duplicates
-
-        # cut vectors if a specific dimension is specified
-        if matryoshka_vector_dimension is not None:
-            sentence_pb.vector.extend(vector.vector[:matryoshka_vector_dimension])
-        else:
-            sentence_pb.vector.extend(vector.vector)
+        sentence_pb.vector.extend(vector.vector[:vector_dimension])
 
         # we only care about start/stop position of the paragraph for a given sentence here
         # the key has the sentence position

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/knowledgebox.py

@@ -57,6 +57,7 @@ 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
@@ -103,9 +104,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,
@@ -120,7 +121,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]]] = []
@@ -523,11 +524,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)