nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/ingest/orm/knowledgebox.py

@@ -17,14 +17,15 @@
 # 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 collections.abc import AsyncGenerator, Callable, Coroutine, Sequence
 from datetime import datetime
 from functools import partial
-from typing import Any, AsyncGenerator, Callable, Coroutine, Optional, Sequence
+from typing import Any
 from uuid import uuid4
 
 from grpc import StatusCode
 from grpc.aio import AioRpcError
-from nidx_protos import noderesources_pb2
+from nidx_protos import nidx_pb2, noderesources_pb2
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardNotFound
@@ -36,7 +37,6 @@ from nucliadb.common.datamanagers.resources import (
     KB_RESOURCE_SLUG_BASE,
 )
 from nucliadb.common.external_index_providers.base import VectorsetExternalIndex
-from nucliadb.common.external_index_providers.pinecone import PineconeIndexManager
 from nucliadb.common.maindb.driver import Driver, Transaction
 from nucliadb.common.maindb.pg import PGTransaction
 from nucliadb.common.nidx import get_nidx_api_client
@@ -53,7 +53,6 @@ from nucliadb.migrator.utils import get_latest_version
 from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_protos.knowledgebox_pb2 import (
     CreateExternalIndexProviderMetadata,
-    ExternalIndexProviderType,
     KnowledgeBoxConfig,
     SemanticModelMetadata,
     StoredExternalIndexProviderMetadata,
@@ -90,7 +89,7 @@ class KnowledgeBox:
         self.txn = txn
         self.storage = storage
         self.kbid = kbid
-        self._config: Optional[KnowledgeBoxConfig] = None
+        self._config: KnowledgeBoxConfig | None = None
 
     @staticmethod
     def new_unique_kbid() -> str:
@@ -110,6 +109,7 @@ class KnowledgeBox:
         external_index_provider: CreateExternalIndexProviderMetadata = CreateExternalIndexProviderMetadata(),
         hidden_resources_enabled: bool = False,
         hidden_resources_hide_on_creation: bool = False,
+        prewarm_enabled: bool = False,
     ) -> tuple[str, str]:
         """Creates a new knowledge box and return its id and slug."""
 
@@ -196,6 +196,7 @@ class KnowledgeBox:
                     migration_version=get_latest_version(),
                     hidden_resources_enabled=hidden_resources_enabled,
                     hidden_resources_hide_on_creation=hidden_resources_hide_on_creation,
+                    prewarm_enabled=prewarm_enabled,
                 )
                 config.external_index_provider.CopyFrom(stored_external_index_provider)
                 await datamanagers.kb.set_config(txn, kbid=kbid, config=config)
@@ -222,7 +223,7 @@ class KnowledgeBox:
                 shard_manager = get_shard_manager()
                 # XXX creating a shard is a slow IO operation that requires a write
                 # txn to be open!
-                await shard_manager.create_shard_by_kbid(txn, kbid)
+                await shard_manager.create_shard_by_kbid(txn, kbid, prewarm_enabled=prewarm_enabled)
                 # shards don't need a rollback as they will be eventually purged
 
                 await txn.commit()
@@ -245,39 +246,86 @@ class KnowledgeBox:
     @classmethod
     async def update(
         cls,
-        txn: Transaction,
-        uuid: str,
-        slug: Optional[str] = None,
-        config: Optional[KnowledgeBoxConfig] = None,
+        driver: Driver,
+        kbid: str,
+        *,
+        slug: str | None = None,
+        title: str | None = None,
+        description: str | None = None,
+        migration_version: int | None = None,
+        external_index_provider: StoredExternalIndexProviderMetadata | None = None,
+        hidden_resources_enabled: bool | None = None,
+        hidden_resources_hide_on_creation: bool | None = None,
+        prewarm_enabled: bool | None = None,
     ) -> str:
-        exist = await datamanagers.kb.get_config(txn, kbid=uuid, for_update=True)
-        if not exist:
-            raise datamanagers.exceptions.KnowledgeBoxNotFound()
-
-        if slug:
-            await txn.delete(datamanagers.kb.KB_SLUGS.format(slug=exist.slug))
-            await txn.set(
-                datamanagers.kb.KB_SLUGS.format(slug=slug),
-                uuid.encode(),
-            )
-            if config:
-                config.slug = slug
-            else:
-                exist.slug = slug
+        async with driver.rw_transaction() as txn:
+            stored = await datamanagers.kb.get_config(txn, kbid=kbid, for_update=True)
+            if not stored:
+                raise datamanagers.exceptions.KnowledgeBoxNotFound()
+
+            if slug:
+                await txn.delete(datamanagers.kb.KB_SLUGS.format(slug=stored.slug))
+                await txn.set(
+                    datamanagers.kb.KB_SLUGS.format(slug=slug),
+                    kbid.encode(),
+                )
+                stored.slug = slug
 
-        if config and exist != config:
-            exist.MergeFrom(config)
-            exist.hidden_resources_enabled = config.hidden_resources_enabled
-            exist.hidden_resources_hide_on_creation = config.hidden_resources_hide_on_creation
+            if title is not None:
+                stored.title = title
+            if description is not None:
+                stored.description = description
 
-        if exist.hidden_resources_hide_on_creation and not exist.hidden_resources_enabled:
-            raise KnowledgeBoxCreationError(
-                "Cannot hide new resources if the hidden resources feature is disabled"
-            )
+            if migration_version is not None:
+                stored.migration_version = migration_version
+
+            if external_index_provider is not None:
+                stored.external_index_provider.MergeFrom(external_index_provider)
+
+            if hidden_resources_enabled is not None:
+                stored.hidden_resources_enabled = hidden_resources_enabled
+            if hidden_resources_hide_on_creation is not None:
+                stored.hidden_resources_hide_on_creation = hidden_resources_hide_on_creation
+
+            update_nidx_prewarm = None
+            if prewarm_enabled is not None:
+                if stored.prewarm_enabled != prewarm_enabled:
+                    update_nidx_prewarm = prewarm_enabled
+                stored.prewarm_enabled = prewarm_enabled
+
+            if stored.hidden_resources_hide_on_creation and not stored.hidden_resources_enabled:
+                raise KnowledgeBoxCreationError(
+                    "Cannot hide new resources if the hidden resources feature is disabled"
+                )
+
+            await datamanagers.kb.set_config(txn, kbid=kbid, config=stored)
+
+            await txn.commit()
+
+        if update_nidx_prewarm is not None:
+            await cls.configure_shards(driver, kbid, prewarm=update_nidx_prewarm)
 
-        await datamanagers.kb.set_config(txn, kbid=uuid, config=exist)
+        return kbid
+
+    @classmethod
+    async def configure_shards(cls, driver: Driver, kbid: str, *, prewarm: bool):
+        shards_obj = await datamanagers.atomic.cluster.get_kb_shards(kbid=kbid)
+        if shards_obj is None:
+            logger.warning(f"Shards not found for KB while updating pre-warm flag", extra={"kbid": kbid})
+            return
 
-        return uuid
+        nidx_shard_ids = [shard.nidx_shard_id for shard in shards_obj.shards]
+
+        nidx_api = get_nidx_api_client()
+        if nidx_api is not None and len(nidx_shard_ids) > 0:
+            configs = [
+                nidx_pb2.ShardConfig(
+                    shard_id=shard_id,
+                    prewarm_enabled=prewarm,
+                )
+                for shard_id in nidx_shard_ids
+            ]
+            await nidx_api.ConfigureShards(nidx_pb2.ShardsConfig(configs=configs))
 
     @classmethod
     async def delete(cls, driver: Driver, kbid: str):
@@ -381,7 +429,7 @@ class KnowledgeBox:
             await txn.delete_by_prefix(prefix)
             await txn.commit()
 
-    async def get_resource_shard(self, shard_id: str) -> Optional[writer_pb2.ShardObject]:
+    async def get_resource_shard(self, shard_id: str) -> writer_pb2.ShardObject | None:
         async with datamanagers.with_ro_transaction() as txn:
             pb = await datamanagers.cluster.get_kb_shards(txn, kbid=self.kbid)
             if pb is None:
@@ -392,18 +440,8 @@ class KnowledgeBox:
                 return shard
         return None
 
-    async def get(self, uuid: str) -> Optional[Resource]:
-        basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kbid, rid=uuid)
-        if basic is None:
-            return None
-        return Resource(
-            txn=self.txn,
-            storage=self.storage,
-            kb=self,
-            uuid=uuid,
-            basic=basic,
-            disable_vectors=False,
-        )
+    async def get(self, uuid: str) -> Resource | None:
+        return await Resource.get(self.txn, self.kbid, uuid)
 
     async def maindb_delete_resource(self, uuid: str):
         basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kbid, rid=uuid)
@@ -432,7 +470,7 @@ class KnowledgeBox:
         with processor_observer({"type": "delete_resource_storage"}):
             await self.storage_delete_resource(uuid)
 
-    async def get_resource_uuid_by_slug(self, slug: str) -> Optional[str]:
+    async def get_resource_uuid_by_slug(self, slug: str) -> str | None:
         return await datamanagers.resources.get_resource_uuid_from_slug(
             self.txn, kbid=self.kbid, slug=slug
         )
@@ -449,7 +487,7 @@ class KnowledgeBox:
                 key_ok = True
         return slug
 
-    async def add_resource(self, uuid: str, slug: str, basic: Optional[Basic] = None) -> Resource:
+    async def add_resource(self, uuid: str, slug: str, basic: Basic | None = None) -> Resource:
         if basic is None:
             basic = Basic()
         if slug == "":
@@ -461,7 +499,7 @@ class KnowledgeBox:
         return Resource(
             storage=self.storage,
             txn=self.txn,
-            kb=self,
+            kbid=self.kbid,
             uuid=uuid,
             basic=basic,
             disable_vectors=False,
@@ -476,7 +514,7 @@ class KnowledgeBox:
                 yield Resource(
                     self.txn,
                     self.storage,
-                    self,
+                    self.kbid,
                     uuid,
                     disable_vectors=False,
                 )
@@ -535,10 +573,7 @@ class KnowledgeBox:
         request: CreateExternalIndexProviderMetadata,
         indexes: list[VectorsetExternalIndex],
     ) -> StoredExternalIndexProviderMetadata:
-        if request.type != ExternalIndexProviderType.PINECONE:
-            return StoredExternalIndexProviderMetadata(type=request.type)
-        # Only pinecone is supported for now
-        return await PineconeIndexManager.create_indexes(kbid, request, indexes)
+        return StoredExternalIndexProviderMetadata(type=request.type)
 
     @classmethod
     async def _maybe_delete_external_indexes(
@@ -546,10 +581,7 @@ class KnowledgeBox:
         kbid: str,
         stored: StoredExternalIndexProviderMetadata,
     ) -> None:
-        if stored.type != ExternalIndexProviderType.PINECONE:
-            return
-        # Only pinecone is supported for now
-        await PineconeIndexManager.delete_indexes(kbid, stored)
+        return
 
 
 def chunker(seq: Sequence, size: int):

nucliadb/ingest/orm/processor/auditing.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from nucliadb.common.maindb.driver import Driver
-from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import audit_pb2, writer_pb2
 from nucliadb_protos.resources_pb2 import FieldType
@@ -35,8 +34,7 @@ async def collect_audit_fields(
 
     audit_storage_fields: list[audit_pb2.AuditField] = []
     async with driver.ro_transaction() as txn:
-        kb = KnowledgeBox(txn, storage, message.kbid)
-        resource = Resource(txn, storage, kb, message.uuid)
+        resource = Resource(txn, storage, message.kbid, message.uuid)
         field_keys = await resource.get_fields_ids()
 
         for field_id, field_type in iterate_auditable_fields(field_keys, message):

nucliadb/ingest/orm/processor/data_augmentation.py

@@ -20,7 +20,6 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Optional
 
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb.ingest.processing import ProcessingEngine
@@ -94,7 +93,7 @@ def _generate_processing_payload_for_fields(
     rid: str,
     fields: GeneratedFields,
     bm: writer_pb2.BrokerMessage,
-) -> Optional[PushPayload]:
+) -> PushPayload | None:
     partitioning = get_partitioning()
     partition = partitioning.generate_partition(kbid, rid)