nucliadb 6.3.7.post4116__py3-none-any.whl → 6.3.7.post4119__py3-none-any.whl

migrations/0017_multiple_writable_shards.py

@@ -31,7 +31,6 @@ future multiple writable shards will be possible.
 
 import logging
 
-from nucliadb.common import datamanagers
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -41,18 +40,22 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    async with context.kv_driver.transaction() as txn:
-        shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
-        if shards is None:
-            logger.error("KB without shards", extra={"kbid": kbid})
-            return
+    pass
 
-        for shard_object in shards.shards:
-            shard_object.read_only = True
-        shards.shards[shards.actual].read_only = False
+    # No longer relevant with nidx
 
-        # just ensure we're writing it correctly
-        assert [shard_object.read_only for shard_object in shards.shards].count(False) == 1
+    # async with context.kv_driver.transaction() as txn:
+    #     shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
+    #     if shards is None:
+    #         logger.error("KB without shards", extra={"kbid": kbid})
+    #         return
 
-        await datamanagers.cluster.update_kb_shards(txn, kbid=kbid, shards=shards)
-        await txn.commit()
+    #     for shard_object in shards.shards:
+    #         shard_object.read_only = True
+    #     shards.shards[shards.actual].read_only = False
+
+    #     # just ensure we're writing it correctly
+    #     assert [shard_object.read_only for shard_object in shards.shards].count(False) == 1
+
+    #     await datamanagers.cluster.update_kb_shards(txn, kbid=kbid, shards=shards)
+    #     await txn.commit()

migrations/0025_assign_models_to_kbs_v2.py

@@ -38,10 +38,7 @@ import logging
 from nucliadb import learning_proxy
 from nucliadb.common import datamanagers
 from nucliadb.migrator.context import ExecutionContext
-from nucliadb_protos import (
-    knowledgebox_pb2,
-    nodewriter_pb2,
-)
+from nucliadb_protos import knowledgebox_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -97,10 +94,10 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
 
     default_vectorset = knowledgebox_pb2.VectorSetConfig(
         vectorset_id=vectorset_id,
-        vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
+        vectorset_index_config=knowledgebox_pb2.VectorIndexConfig(
             vector_dimension=maindb_vector_dimension,
             similarity=maindb_similarity,
-            vector_type=nodewriter_pb2.VectorType.DENSE_F32,  # we only support this for now
+            vector_type=knowledgebox_pb2.VectorType.DENSE_F32,  # we only support this for now
             normalize_vectors=maindb_normalize_vectors,
         ),
         matryoshka_dimensions=maindb_matryoshka_dimensions,

nucliadb/common/cluster/base.py

@@ -20,15 +20,17 @@
 from abc import ABCMeta, abstractmethod
 from typing import AsyncIterator
 
-from nucliadb_protos import nodereader_pb2, noderesources_pb2, utils_pb2
-from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
-from nucliadb_protos.nodewriter_pb2 import (
+from nidx_protos import nodereader_pb2, noderesources_pb2
+from nidx_protos.nodereader_pb2_grpc import NodeReaderStub
+from nidx_protos.nodewriter_pb2 import (
     NewShardRequest,
     NewVectorSetRequest,
     OpStatus,
     VectorIndexConfig,
 )
-from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
+from nidx_protos.nodewriter_pb2_grpc import NodeWriterStub
+
+from nucliadb_protos import utils_pb2
 
 
 class AbstractIndexNode(metaclass=ABCMeta):

nucliadb/common/cluster/grpc_node_dummy.py

@@ -19,19 +19,20 @@
 #
 from typing import Any
 
-from nucliadb_protos.nodereader_pb2 import (
+from nidx_protos.nodereader_pb2 import (
     EdgeList,
     RelationEdge,
 )
-from nucliadb_protos.noderesources_pb2 import (
+from nidx_protos.noderesources_pb2 import (
     EmptyResponse,
     ShardCreated,
     ShardId,
     ShardIds,
     VectorSetList,
 )
-from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
-from nucliadb_protos.nodewriter_pb2 import OpStatus
+from nidx_protos.noderesources_pb2 import Shard as NodeResourcesShard
+from nidx_protos.nodewriter_pb2 import OpStatus
+
 from nucliadb_protos.utils_pb2 import Relation
 
 

nucliadb/common/cluster/manager.py

@@ -22,6 +22,9 @@ import logging
 import uuid
 from typing import Any, Awaitable, Callable, Optional
 
+from nidx_protos import noderesources_pb2, nodewriter_pb2
+from nidx_protos.nodewriter_pb2 import IndexMessage, IndexMessageSource, NewShardRequest, TypeMessage
+
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import (
@@ -31,13 +34,8 @@ from nucliadb.common.cluster.exceptions import (
 )
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.common.nidx import get_nidx, get_nidx_api_client, get_nidx_fake_node
-from nucliadb_protos import (
-    knowledgebox_pb2,
-    noderesources_pb2,
-    nodewriter_pb2,
-    writer_pb2,
-)
-from nucliadb_protos.nodewriter_pb2 import IndexMessage, IndexMessageSource, NewShardRequest, TypeMessage
+from nucliadb.common.vector_index_config import nucliadb_index_config_to_nidx
+from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_telemetry import errors
 from nucliadb_utils.utilities import get_storage
 
@@ -123,7 +121,7 @@ class KBShardManager:
             raise ShardsNotFound(msg)
 
         vectorsets = {
-            vectorset_id: vectorset_config.vectorset_index_config
+            vectorset_id: nucliadb_index_config_to_nidx(vectorset_config.vectorset_index_config)
             async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(txn, kbid=kbid)
         }
 
@@ -256,7 +254,7 @@ class KBShardManager:
 
         async def _create_vectorset(node: AbstractIndexNode, shard_id: str):
             vectorset_id = config.vectorset_id
-            index_config = config.vectorset_index_config
+            index_config = nucliadb_index_config_to_nidx(config.vectorset_index_config)
             result = await node.add_vectorset(shard_id, vectorset_id, index_config)
             if result.status != result.Status.OK:
                 raise NodeError(

nucliadb/common/cluster/rebalance.py

@@ -20,11 +20,12 @@
 import asyncio
 import logging
 
+from nidx_protos import nodereader_pb2, noderesources_pb2
+
 from nucliadb.common import datamanagers, locking
 from nucliadb.common.cluster.manager import choose_node
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.context import ApplicationContext
-from nucliadb_protos import nodereader_pb2, noderesources_pb2
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
 from nucliadb_telemetry.utils import setup_telemetry

nucliadb/common/cluster/rollover.py

@@ -31,6 +31,7 @@ from nucliadb.common.external_index_providers.manager import (
     get_external_index_manager,
 )
 from nucliadb.common.nidx import get_nidx_fake_node
+from nucliadb.common.vector_index_config import nucliadb_index_config_to_nidx
 from nucliadb.migrator.settings import settings
 from nucliadb_protos import writer_pb2
 from nucliadb_telemetry import errors
@@ -137,9 +138,8 @@ async def create_rollover_shards(
     created_shards = []
     try:
         for shard in kb_shards.shards:
-            shard.ClearField("replicas")
             vectorsets = {
-                vectorset_id: vectorset_config.vectorset_index_config
+                vectorset_id: nucliadb_index_config_to_nidx(vectorset_config.vectorset_index_config)
                 async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(txn, kbid=kbid)
             }
 

nucliadb/common/cluster/utils.py

@@ -21,6 +21,7 @@ import logging
 from typing import TYPE_CHECKING, Optional, Union
 
 import backoff
+from nidx_protos import nodereader_pb2
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.manager import (
@@ -30,7 +31,7 @@ from nucliadb.common.cluster.manager import (
 from nucliadb.common.cluster.settings import settings
 from nucliadb.ingest.orm import index_message
 from nucliadb.ingest.orm.resource import Resource
-from nucliadb_protos import nodereader_pb2, writer_pb2
+from nucliadb_protos import writer_pb2
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 if TYPE_CHECKING:  # pragma: no cover

nucliadb/common/datamanagers/rollover.py

@@ -79,9 +79,8 @@ async def is_rollover_shard(txn: Transaction, *, kbid: str, shard_id: str) -> bo
         return False
 
     for shard_obj in shards.shards:
-        for replica_obj in shard_obj.replicas:
-            if shard_id == replica_obj.shard.id:
-                return True
+        if shard_id == shard_obj.nidx_shard_id:
+            return True
     return False
 
 

nucliadb/common/external_index_providers/base.py

@@ -22,6 +22,8 @@ import logging
 from dataclasses import dataclass
 from typing import Any, Iterator, Optional
 
+from nidx_protos.nodereader_pb2 import SearchRequest
+from nidx_protos.noderesources_pb2 import Resource
 from pydantic import BaseModel
 
 from nucliadb.common.counters import IndexCounts
@@ -33,8 +35,6 @@ 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
 

nucliadb/common/external_index_providers/pinecone.py

@@ -25,6 +25,8 @@ from uuid import uuid4
 
 import backoff
 from cachetools import TTLCache
+from nidx_protos.nodereader_pb2 import FilterExpression, SearchRequest
+from nidx_protos.noderesources_pb2 import IndexParagraph, Resource, VectorSentence
 from pydantic import BaseModel
 
 from nucliadb.common.counters import IndexCounts
@@ -40,8 +42,6 @@ from nucliadb.common.ids import ParagraphId, VectorId
 from nucliadb_models.search import SCORE_TYPE, TextPosition
 from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import utils_pb2
-from nucliadb_protos.nodereader_pb2 import FilterExpression, SearchRequest
-from nucliadb_protos.noderesources_pb2 import IndexParagraph, Resource, VectorSentence
 from nucliadb_telemetry.metrics import Observer
 from nucliadb_utils.aiopynecone.client import DataPlane, FilterOperator, LogicalOperator
 from nucliadb_utils.aiopynecone.exceptions import (

nucliadb/common/nidx.py

@@ -22,14 +22,14 @@ import os
 from typing import Optional
 
 from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxIndexerStub, NidxSearcherStub
+from nidx_protos.nodewriter_pb2 import (
+    IndexMessage,
+)
 
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.settings import settings
 from nucliadb.ingest.settings import DriverConfig
 from nucliadb.ingest.settings import settings as ingest_settings
-from nucliadb_protos.nodewriter_pb2 import (
-    IndexMessage,
-)
 from nucliadb_utils import logger
 from nucliadb_utils.grpc import get_traced_grpc_channel
 from nucliadb_utils.nats import NatsConnectionManager

nucliadb/common/vector_index_config.py

@@ -0,0 +1,39 @@
+# 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 nidx_protos import nodewriter_pb2 as Nidx
+
+from nucliadb_protos import knowledgebox_pb2 as Nucliadb
+
+
+def nucliadb_vector_type_to_nidx(nucliadb: Nucliadb.VectorType.ValueType) -> Nidx.VectorType.ValueType:
+    if nucliadb == Nucliadb.DENSE_F32:
+        return Nidx.DENSE_F32
+    else:  # pragma: nocover
+        raise Exception("Unknown vector type")
+
+
+def nucliadb_index_config_to_nidx(nucliadb: Nucliadb.VectorIndexConfig) -> Nidx.VectorIndexConfig:
+    return Nidx.VectorIndexConfig(
+        normalize_vectors=nucliadb.normalize_vectors,
+        similarity=nucliadb.similarity,
+        vector_dimension=nucliadb.vector_dimension,
+        vector_type=nucliadb_vector_type_to_nidx(nucliadb.vector_type),
+    )

nucliadb/ingest/consumer/auditing.py

@@ -23,12 +23,14 @@ import logging
 import uuid
 from functools import partial
 
+from nidx_protos import nodereader_pb2, noderesources_pb2
+
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardsNotFound
 from nucliadb.common.cluster.manager import choose_node
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.constants import AVG_PARAGRAPH_SIZE_BYTES
-from nucliadb_protos import audit_pb2, nodereader_pb2, noderesources_pb2, writer_pb2
+from nucliadb_protos import audit_pb2, writer_pb2
 from nucliadb_utils import const
 from nucliadb_utils.audit.audit import AuditStorage
 from nucliadb_utils.cache.pubsub import PubSubDriver

nucliadb/ingest/consumer/shard_creator.py

@@ -22,11 +22,13 @@ import logging
 import uuid
 from functools import partial
 
+from nidx_protos import nodereader_pb2, noderesources_pb2
+
 from nucliadb.common import locking
 from nucliadb.common.cluster.manager import choose_node
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.maindb.driver import Driver
-from nucliadb_protos import nodereader_pb2, noderesources_pb2, writer_pb2
+from nucliadb_protos import writer_pb2
 from nucliadb_utils import const
 from nucliadb_utils.cache.pubsub import PubSubDriver
 from nucliadb_utils.storages.storage import Storage

nucliadb/ingest/orm/brain.py

@@ -22,21 +22,22 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional
 
+from nidx_protos.noderesources_pb2 import IndexParagraph as BrainParagraph
+from nidx_protos.noderesources_pb2 import (
+    IndexRelation,
+    ParagraphMetadata,
+    Representation,
+    ResourceID,
+)
+from nidx_protos.noderesources_pb2 import Position as TextPosition
+from nidx_protos.noderesources_pb2 import Resource as PBBrainResource
+
 from nucliadb.common import ids
 from nucliadb.ingest import logger
 from nucliadb.ingest.orm.utils import compute_paragraph_key
 from nucliadb_models.labels import BASE_LABELS, LABEL_HIDDEN, flatten_resource_labels
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_protos import utils_pb2
-from nucliadb_protos.noderesources_pb2 import IndexParagraph as BrainParagraph
-from nucliadb_protos.noderesources_pb2 import (
-    IndexRelation,
-    ParagraphMetadata,
-    Representation,
-    ResourceID,
-)
-from nucliadb_protos.noderesources_pb2 import Position as TextPosition
-from nucliadb_protos.noderesources_pb2 import Resource as PBBrainResource
 from nucliadb_protos.resources_pb2 import (
     Basic,
     ExtractedText,

nucliadb/ingest/orm/brain_v2.py

@@ -22,6 +22,16 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional
 
+from nidx_protos.noderesources_pb2 import IndexParagraph as BrainParagraph
+from nidx_protos.noderesources_pb2 import (
+    IndexRelation,
+    ParagraphMetadata,
+    Representation,
+    ResourceID,
+)
+from nidx_protos.noderesources_pb2 import Position as TextPosition
+from nidx_protos.noderesources_pb2 import Resource as PBBrainResource
+
 from nucliadb.common import ids
 from nucliadb.ingest import logger
 from nucliadb.ingest.orm.metrics import brain_observer as observer
@@ -29,15 +39,6 @@ from nucliadb.ingest.orm.utils import compute_paragraph_key
 from nucliadb_models.labels import BASE_LABELS, LABEL_HIDDEN, flatten_resource_labels
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_protos import utils_pb2
-from nucliadb_protos.noderesources_pb2 import IndexParagraph as BrainParagraph
-from nucliadb_protos.noderesources_pb2 import (
-    IndexRelation,
-    ParagraphMetadata,
-    Representation,
-    ResourceID,
-)
-from nucliadb_protos.noderesources_pb2 import Position as TextPosition
-from nucliadb_protos.noderesources_pb2 import Resource as PBBrainResource
 from nucliadb_protos.resources_pb2 import (
     Basic,
     ExtractedText,

nucliadb/ingest/orm/entities.py

@@ -21,6 +21,14 @@
 import asyncio
 from typing import AsyncGenerator, Optional
 
+from nidx_protos.nodereader_pb2 import (
+    Faceted,
+    GraphSearchRequest,
+    GraphSearchResponse,
+    SearchRequest,
+    SearchResponse,
+)
+
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import (
@@ -43,13 +51,6 @@ from nucliadb_protos.knowledgebox_pb2 import (
     EntitiesGroupSummary,
     Entity,
 )
-from nucliadb_protos.nodereader_pb2 import (
-    Faceted,
-    GraphSearchRequest,
-    GraphSearchResponse,
-    SearchRequest,
-    SearchResponse,
-)
 from nucliadb_protos.utils_pb2 import RelationNode
 from nucliadb_protos.writer_pb2 import GetEntitiesResponse
 

nucliadb/ingest/orm/index_message.py

@@ -22,6 +22,8 @@
 import asyncio
 from typing import Optional
 
+from nidx_protos.noderesources_pb2 import Resource as IndexMessage
+
 from nucliadb.common import datamanagers
 from nucliadb.ingest.fields.exceptions import FieldAuthorNotFound
 from nucliadb.ingest.fields.file import File
@@ -29,7 +31,6 @@ from nucliadb.ingest.orm.brain_v2 import ResourceBrainV2 as ResourceBrain
 from nucliadb.ingest.orm.metrics import index_message_observer as observer
 from nucliadb.ingest.orm.resource import Resource, get_file_page_positions
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
-from nucliadb_protos.noderesources_pb2 import Resource as IndexMessage
 from nucliadb_protos.resources_pb2 import Basic, FieldID, FieldType
 from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils import const

nucliadb/ingest/orm/knowledgebox.py

@@ -24,6 +24,7 @@ from uuid import uuid4
 
 from grpc import StatusCode
 from grpc.aio import AioRpcError
+from nidx_protos import noderesources_pb2
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardNotFound
@@ -49,7 +50,7 @@ from nucliadb.ingest.orm.metrics import processor_observer
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb.ingest.orm.utils import choose_matryoshka_dimension, compute_paragraph_key
 from nucliadb.migrator.utils import get_latest_version
-from nucliadb_protos import knowledgebox_pb2, noderesources_pb2, nodewriter_pb2, writer_pb2
+from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_protos.knowledgebox_pb2 import (
     CreateExternalIndexProviderMetadata,
     ExternalIndexProviderType,
@@ -165,10 +166,10 @@ class KnowledgeBox:
 
                     vectorset_config = knowledgebox_pb2.VectorSetConfig(
                         vectorset_id=vectorset_id,
-                        vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
+                        vectorset_index_config=knowledgebox_pb2.VectorIndexConfig(
                             similarity=semantic_model.similarity_function,
                             # XXX: hardcoded value
-                            vector_type=nodewriter_pb2.VectorType.DENSE_F32,
+                            vector_type=knowledgebox_pb2.VectorType.DENSE_F32,
                             normalize_vectors=len(semantic_model.matryoshka_dimensions) > 0,
                             vector_dimension=dimension,
                         ),

nucliadb/ingest/orm/processor/pgcatalog.py

@@ -20,10 +20,11 @@
 
 from typing import cast
 
+from nidx_protos.noderesources_pb2 import Resource as IndexMessage
+
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.common.maindb.pg import PGDriver, PGTransaction
 from nucliadb.common.maindb.utils import get_driver
-from nucliadb_protos.noderesources_pb2 import Resource as IndexMessage
 from nucliadb_telemetry import metrics
 
 from ..resource import Resource

nucliadb/ingest/orm/processor/processor.py

@@ -24,6 +24,8 @@ from typing import Optional
 import aiohttp.client_exceptions
 import nats.errors
 import nats.js.errors
+from nidx_protos import noderesources_pb2, nodewriter_pb2
+from nidx_protos.noderesources_pb2 import Resource as PBBrainResource
 
 from nucliadb.common import datamanagers, locking
 from nucliadb.common.cluster.settings import settings as cluster_settings
@@ -50,12 +52,9 @@ from nucliadb.ingest.orm.processor.data_augmentation import (
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import (
     knowledgebox_pb2,
-    noderesources_pb2,
-    nodewriter_pb2,
     resources_pb2,
     writer_pb2,
 )
-from nucliadb_protos.noderesources_pb2 import Resource as PBBrainResource
 from nucliadb_telemetry import errors
 from nucliadb_utils import const
 from nucliadb_utils.cache.pubsub import PubSubDriver

nucliadb/metrics_exporter.py

@@ -22,6 +22,8 @@ from __future__ import annotations
 import asyncio
 from typing import AsyncGenerator, Callable, Tuple, cast
 
+from nidx_protos.noderesources_pb2 import EmptyQuery, NodeMetadata
+
 from nucliadb import logger
 from nucliadb.common import datamanagers
 from nucliadb.common.context import ApplicationContext
@@ -29,7 +31,6 @@ from nucliadb.common.maindb.pg import PGDriver
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.common.nidx import get_nidx_api_client
 from nucliadb.migrator.datamanager import MigrationsDataManager
-from nucliadb_protos.noderesources_pb2 import EmptyQuery, NodeMetadata
 from nucliadb_telemetry import metrics
 from nucliadb_telemetry.logs import setup_logging
 from nucliadb_telemetry.utils import setup_telemetry

nucliadb/search/api/v1/knowledgebox.py

@@ -24,6 +24,7 @@ from fastapi import HTTPException, Request
 from fastapi_versioning import version
 from grpc import StatusCode as GrpcStatusCode
 from grpc.aio import AioRpcError
+from nidx_protos.noderesources_pb2 import Shard
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardsNotFound
@@ -44,7 +45,6 @@ from nucliadb_models.search import (
     KnowledgeboxCounters,
     SearchParamDefaults,
 )
-from nucliadb_protos.noderesources_pb2 import Shard
 from nucliadb_protos.writer_pb2 import ShardObject as PBShardObject
 from nucliadb_protos.writer_pb2 import Shards
 from nucliadb_telemetry import errors

nucliadb/search/requesters/utils.py

@@ -26,6 +26,14 @@ from fastapi import HTTPException
 from google.protobuf.json_format import MessageToDict
 from grpc import StatusCode as GrpcStatusCode
 from grpc.aio import AioRpcError
+from nidx_protos.nodereader_pb2 import (
+    GraphSearchRequest,
+    GraphSearchResponse,
+    SearchRequest,
+    SearchResponse,
+    SuggestRequest,
+    SuggestResponse,
+)
 
 from nucliadb.common.cluster import manager as cluster_manager
 from nucliadb.common.cluster.base import AbstractIndexNode
@@ -38,14 +46,6 @@ from nucliadb.search.search.shards import (
     suggest_shard,
 )
 from nucliadb.search.settings import settings
-from nucliadb_protos.nodereader_pb2 import (
-    GraphSearchRequest,
-    GraphSearchResponse,
-    SearchRequest,
-    SearchResponse,
-    SuggestRequest,
-    SuggestResponse,
-)
 from nucliadb_protos.writer_pb2 import ShardObject as PBShardObject
 from nucliadb_telemetry import errors
 

nucliadb/search/search/chat/query.py

@@ -20,6 +20,11 @@
 import asyncio
 from typing import Iterable, Optional, Union
 
+from nidx_protos.nodereader_pb2 import (
+    GraphSearchResponse,
+    SearchResponse,
+)
+
 from nucliadb.common.models_utils import to_proto
 from nucliadb.search import logger
 from nucliadb.search.predict import AnswerStatusCode
@@ -52,10 +57,6 @@ from nucliadb_models.search import (
     parse_rephrase_prompt,
 )
 from nucliadb_protos import audit_pb2
-from nucliadb_protos.nodereader_pb2 import (
-    GraphSearchResponse,
-    SearchResponse,
-)
 from nucliadb_protos.utils_pb2 import RelationNode
 from nucliadb_telemetry.errors import capture_exception
 from nucliadb_utils.utilities import get_audit

nucliadb/search/search/fetch.py

@@ -21,6 +21,8 @@ import asyncio
 from contextvars import ContextVar
 from typing import Optional
 
+from nidx_protos.nodereader_pb2 import DocumentResult, ParagraphResult
+
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
@@ -30,7 +32,6 @@ from nucliadb.search.search import cache
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import ResourceProperties
-from nucliadb_protos.nodereader_pb2 import DocumentResult, ParagraphResult
 from nucliadb_protos.resources_pb2 import Paragraph
 from nucliadb_utils import const
 from nucliadb_utils.utilities import has_feature