nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/datamanagers/vectorsets.py

@@ -0,0 +1,110 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 AsyncIterator, Optional
+
+from nucliadb.common.datamanagers.utils import get_kv_pb
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb_protos import knowledgebox_pb2
+
+KB_VECTORSETS = "/kbs/{kbid}/vectorsets"
+
+
+class BrokenInvariant(Exception):
+    pass
+
+
+async def initialize(txn: Transaction, *, kbid: str):
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, knowledgebox_pb2.KnowledgeBoxVectorSetsConfig().SerializeToString())
+
+
+async def get(
+    txn: Transaction, *, kbid: str, vectorset_id: str
+) -> Optional[knowledgebox_pb2.VectorSetConfig]:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    index = _find_vectorset(kb_vectorsets, vectorset_id)
+    if index is None:
+        return None
+    return kb_vectorsets.vectorsets[index]
+
+
+async def exists(txn, *, kbid: str, vectorset_id: str) -> bool:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    return _find_vectorset(kb_vectorsets, vectorset_id) is not None
+
+
+async def iter(
+    txn: Transaction, *, kbid: str
+) -> AsyncIterator[tuple[str, knowledgebox_pb2.VectorSetConfig]]:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    for config in kb_vectorsets.vectorsets:
+        yield config.vectorset_id, config
+
+
+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)
+    index = _find_vectorset(kb_vectorsets, config.vectorset_id)
+    if index is None:
+        # adding a new vectorset
+        kb_vectorsets.vectorsets.append(config)
+    else:
+        # updating a vectorset
+        kb_vectorsets.vectorsets[index].CopyFrom(config)
+
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, kb_vectorsets.SerializeToString())
+
+
+async def delete(txn: Transaction, *, kbid: str, vectorset_id: str):
+    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
+
+    del kb_vectorsets.vectorsets[index]
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, kb_vectorsets.SerializeToString())
+
+
+# XXX At some point in the vectorset epic, we should make this key mandatory and
+# fail instead of providing a default
+async def _get_or_default(
+    txn: Transaction,
+    *,
+    kbid: str,
+    for_update: bool = True,
+) -> knowledgebox_pb2.KnowledgeBoxVectorSetsConfig:
+    key = KB_VECTORSETS.format(kbid=kbid)
+    stored = await get_kv_pb(
+        txn, key, knowledgebox_pb2.KnowledgeBoxVectorSetsConfig, for_update=for_update
+    )
+    return stored or knowledgebox_pb2.KnowledgeBoxVectorSetsConfig()
+
+
+def _find_vectorset(
+    kb_vectorsets: knowledgebox_pb2.KnowledgeBoxVectorSetsConfig, vectorset_id: str
+) -> Optional[int]:
+    """Return the position of the vectorset in `vectorsets` or `None` if not found."""
+    for idx, vectorset in enumerate(kb_vectorsets.vectorsets):
+        if vectorset.vectorset_id == vectorset_id:
+            return idx
+    return None

nucliadb/common/external_index_providers/base.py

@@ -0,0 +1,257 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+import abc
+import logging
+from dataclasses import dataclass
+from typing import Any, Iterator, Optional
+
+from pydantic import BaseModel
+
+from nucliadb.common.counters import IndexCounts
+from nucliadb.common.external_index_providers.exceptions import ExternalIndexingError
+from nucliadb.common.ids import ParagraphId
+from nucliadb_models.external_index_providers import ExternalIndexProviderType
+from nucliadb_models.search import SCORE_TYPE, TextPosition
+from nucliadb_protos.knowledgebox_pb2 import (
+    CreateExternalIndexProviderMetadata,
+    StoredExternalIndexProviderMetadata,
+)
+from nucliadb_protos.nodereader_pb2 import SearchRequest
+from nucliadb_protos.noderesources_pb2 import Resource
+from nucliadb_protos.utils_pb2 import VectorSimilarity
+from nucliadb_telemetry.metrics import Observer
+
+logger = logging.getLogger(__name__)
+
+manager_observer = Observer("external_index_manager", labels={"operation": "", "provider": ""})
+
+
+@dataclass
+class VectorsetExternalIndex:
+    """
+    Used to indicate to external index managers the required metadata
+    in order to create an external index for each vectorset
+    """
+
+    vectorset_id: str
+    dimension: int
+    similarity: VectorSimilarity.ValueType
+
+
+class TextBlockMatch(BaseModel):
+    """
+    Model a text block/paragraph retrieved from an external index with all the information
+    needed in order to later hydrate retrieval results.
+    """
+
+    paragraph_id: ParagraphId
+    position: TextPosition
+    score: float
+    score_type: SCORE_TYPE
+    order: int
+    page_with_visual: bool = False
+    fuzzy_search: bool
+    is_a_table: bool = False
+    representation_file: Optional[str] = None
+    paragraph_labels: list[str] = []
+    field_labels: list[str] = []
+    text: Optional[str] = None
+
+
+class QueryResults(BaseModel):
+    """
+    Model for the results of a query to an external index provider.
+    Must be subclassed by the specific external index provider.
+    """
+
+    type: ExternalIndexProviderType
+    results: Any
+
+    def iter_matching_text_blocks(self) -> Iterator[TextBlockMatch]:
+        """
+        Iterates over the paragraphs in the results, by decreasing score.
+        This should be implemented by the specific external index provider.
+        """
+        raise NotImplementedError()
+
+
+class ExternalIndexManager(abc.ABC, metaclass=abc.ABCMeta):
+    """
+    Base class for the external index providers. Must be subclassed by the specific external index provider.
+    """
+
+    type: ExternalIndexProviderType
+    supports_rollover: bool = False
+
+    def __init__(self, kbid: str):
+        self.kbid = kbid
+
+    @classmethod
+    @abc.abstractmethod
+    async def create_indexes(
+        cls,
+        kbid: str,
+        create_request: CreateExternalIndexProviderMetadata,
+        indexes: list[VectorsetExternalIndex],
+    ) -> StoredExternalIndexProviderMetadata: ...
+
+    @classmethod
+    @abc.abstractmethod
+    async def delete_indexes(
+        cls,
+        kbid: str,
+        stored: StoredExternalIndexProviderMetadata,
+    ) -> None: ...
+
+    @abc.abstractmethod
+    async def rollover_create_indexes(
+        self, stored: StoredExternalIndexProviderMetadata
+    ) -> StoredExternalIndexProviderMetadata:  # pragma: no cover
+        """
+        Creates the indexes for the rollover process.
+        In the event of an error, it should rollback any left over indexes.
+        Returns a modified version of the stored external index provider metadata with the new indexes for the rollover.
+        """
+        ...
+
+    @abc.abstractmethod
+    async def rollover_cutover_indexes(self) -> None:  # pragma: no cover
+        """
+        Cutover the indexes for the rollover process.
+        After this operation, the new indexes should be used for queries and the old ones should be deleted.
+        """
+        ...
+
+    @classmethod
+    def get_index_name(cls) -> str:  # pragma: no cover
+        """
+        Returns the name of the index in the external index provider.
+        """
+        raise NotImplementedError()
+
+    async def delete_resource(self, resource_uuid: str) -> None:
+        """
+        Deletes a resource from the external index provider.
+        """
+        logger.info(
+            "Deleting resource to external index",
+            extra={
+                "kbid": self.kbid,
+                "rid": resource_uuid,
+                "provider": self.type.value,
+            },
+        )
+        with manager_observer({"operation": "delete_resource", "provider": self.type.value}):
+            await self._delete_resource(resource_uuid)
+
+    async def index_resource(
+        self, resource_uuid: str, resource_data: Resource, to_rollover_indexes: bool = False
+    ) -> None:
+        """
+        Indexes a resource to the external index provider.
+        """
+        if not self.supports_rollover and to_rollover_indexes:
+            logger.info(
+                "Indexing to rollover indexes not supported",
+                extra={
+                    "kbid": self.kbid,
+                    "rid": resource_uuid,
+                    "provider": self.type.value,
+                },
+            )
+            return
+        logger.info(
+            "Indexing resource to external index",
+            extra={
+                "kbid": self.kbid,
+                "rid": resource_uuid,
+                "provider": self.type.value,
+                "rollover": to_rollover_indexes,
+            },
+        )
+        with manager_observer({"operation": "index_resource", "provider": self.type.value}):
+            try:
+                await self._index_resource(
+                    resource_uuid, resource_data, to_rollover_indexes=to_rollover_indexes
+                )
+            except Exception as ex:
+                raise ExternalIndexingError() from ex
+
+    async def get_index_counts(self) -> IndexCounts:
+        """
+        Returns the index counts for the external index provider.
+        """
+        logger.debug(
+            "Getting index counts from external index",
+            extra={
+                "kbid": self.kbid,
+                "provider": self.type.value,
+            },
+        )
+        with manager_observer({"operation": "get_index_counts", "provider": self.type.value}):
+            return await self._get_index_counts()
+
+    async def query(self, request: SearchRequest) -> QueryResults:
+        """
+        Queries the external index provider and returns the results.
+        """
+        logger.info(
+            "Querying external index",
+            extra={
+                "kbid": self.kbid,
+                "provider": self.type.value,
+            },
+        )
+        with manager_observer({"operation": "query", "provider": self.type.value}):
+            return await self._query(request)
+
+    @abc.abstractmethod
+    async def _delete_resource(self, resource_uuid: str) -> None:  # pragma: no cover
+        """
+        Makes sure that all vectors associated with the resource are deleted from the external index provider.
+        """
+        ...
+
+    @abc.abstractmethod
+    async def _index_resource(
+        self, resource_uuid: str, resource_data: Resource, to_rollover_indexes: bool = False
+    ) -> None:  # pragma: no cover
+        """
+        Adapts the Resource (aka brain) to the external index provider's index format and indexes it.
+        Params:
+        - resource_uuid: the resource's UUID
+        - resource_data: the resource index data
+        - to_rollover_indexes: whether to index to the rollover indexes or the main indexes
+        """
+        ...
+
+    @abc.abstractmethod
+    async def _query(self, request: SearchRequest) -> QueryResults:  # pragma: no cover
+        """
+        Adapts the Nucliadb's search request to the external index provider's query format and returns the results.
+        """
+        ...
+
+    @abc.abstractmethod
+    async def _get_index_counts(self) -> IndexCounts:  # pragma: no cover
+        """
+        Returns the index counts for the external index provider.
+        """
+        ...

nucliadb/{ingest/tests/unit/test_cache.py → common/external_index_providers/exceptions.py}

@@ -17,15 +17,16 @@
 # 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 nucliadb.common.cluster.index_node import READ_CONNECTIONS, WRITE_CONNECTIONS
-from nucliadb.ingest.cache import clear_ingest_cache
 
 
-def test_clear_ingest_cache():
-    READ_CONNECTIONS["addr1"] = "conn1"
-    WRITE_CONNECTIONS["addr2"] = "conn2"
+class ExternalIndexCreationError(Exception):
+    def __init__(self, provider: str, message: str):
+        self.provider = provider
+        self.message = message
+        super().__init__(f"{provider} index creation error: {message}")
 
-    clear_ingest_cache()
 
-    assert len(READ_CONNECTIONS) == 0
-    assert len(WRITE_CONNECTIONS) == 0
+class ExternalIndexingError(Exception):
+    """
+    Raised when an error occurs while indexing a resource in an external index.
+    """

nucliadb/common/external_index_providers/manager.py

@@ -0,0 +1,101 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 async_lru
+
+from nucliadb.common import datamanagers
+from nucliadb.common.external_index_providers.base import ExternalIndexManager
+from nucliadb.common.external_index_providers.pinecone import PineconeIndexManager
+from nucliadb.common.external_index_providers.settings import settings
+from nucliadb_protos.knowledgebox_pb2 import (
+    ExternalIndexProviderType,
+    StoredExternalIndexProviderMetadata,
+)
+from nucliadb_utils.utilities import get_endecryptor
+
+
+async def get_external_index_manager(
+    kbid: str, for_rollover: bool = False
+) -> Optional[ExternalIndexManager]:
+    """
+    Returns an ExternalIndexManager for the given kbid.
+    If for_rollover is True, the ExternalIndexManager returned will include the rollover indexes (if any).
+    """
+    metadata = await get_external_index_metadata(kbid)
+    if metadata is None or metadata.type != ExternalIndexProviderType.PINECONE:
+        # Only Pinecone is supported for now
+        return None
+
+    api_key = get_endecryptor().decrypt(metadata.pinecone_config.encrypted_api_key)
+    default_vectorset = await get_default_vectorset_id(kbid)
+
+    rollover_indexes = None
+    if for_rollover:
+        rollover_metadata = await get_rollover_external_index_metadata(kbid)
+        if rollover_metadata is not None:
+            rollover_indexes = dict(rollover_metadata.pinecone_config.indexes)
+
+    return PineconeIndexManager(
+        kbid=kbid,
+        api_key=api_key,
+        indexes=dict(metadata.pinecone_config.indexes),
+        upsert_parallelism=settings.pinecone_upsert_parallelism,
+        delete_parallelism=settings.pinecone_delete_parallelism,
+        upsert_timeout=settings.pinecone_upsert_timeout,
+        delete_timeout=settings.pinecone_delete_timeout,
+        default_vectorset=default_vectorset,
+        rollover_indexes=rollover_indexes,
+    )
+
+
+@async_lru.alru_cache(maxsize=None)
+async def get_external_index_metadata(kbid: str) -> Optional[StoredExternalIndexProviderMetadata]:
+    return await datamanagers.atomic.kb.get_external_index_provider_metadata(kbid=kbid)
+
+
+@async_lru.alru_cache(maxsize=None)
+async def get_default_vectorset_id(kbid: str) -> Optional[str]:
+    """
+    While we are transitioning to the new vectorset system, we need to take into account
+    that KBs that have only one semantic model will have the `vectorset_id` field on BrokerMessage.field_vectors
+    set to empty string -- that is the `default` vectorset concept.
+    """
+    async with datamanagers.with_ro_transaction() as txn:
+        vss = []
+        async for vs_id, vs_config in datamanagers.vectorsets.iter(txn, kbid=kbid):
+            vss.append((vs_id, vs_config))
+        if len(vss) == 0:
+            # If there is nothing in the vectorsets key on maindb, we use the "__default__" vectorset as id.
+            return "__default__"
+        if len(vss) == 1:
+            # If there is only one vectorset, return it as the default
+            return vss[0][0]
+        else:
+            # If there are multiple vectorsets, we don't have a default
+            # and we assume the index messages are explicit about the vectorset
+            return None
+
+
+async def get_rollover_external_index_metadata(
+    kbid: str,
+) -> Optional[StoredExternalIndexProviderMetadata]:
+    async with datamanagers.with_ro_transaction() as txn:
+        return await datamanagers.rollover.get_kb_rollover_external_index_metadata(txn, kbid=kbid)