nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/train/generators/token_classifier.py

@@ -21,6 +21,10 @@
 from collections import OrderedDict
 from typing import AsyncGenerator, cast
 
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.train import logger
+from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
     TokenClassificationBatch,
     TokensClassification,
@@ -28,11 +32,6 @@ from nucliadb_protos.dataset_pb2 import (
 )
 from nucliadb_protos.nodereader_pb2 import StreamFilter, StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.ingest.orm.resource import KB_REVERSE
-from nucliadb.train import logger
-from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
-
 NERS_DICT = dict[str, dict[str, list[tuple[int, int]]]]
 POSITION_DICT = OrderedDict[tuple[int, int], tuple[str, str]]
 MAIN = "__main__"
@@ -44,9 +43,7 @@ def token_classification_batch_generator(
     node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TokenClassificationBatch, None]:
-    generator = generate_token_classification_payloads(
-        kbid, trainset, node, shard_replica_id
-    )
+    generator = generate_token_classification_payloads(kbid, trainset, node, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, TokenClassificationBatch)
     return batch_generator
 
@@ -97,21 +94,19 @@ async def get_field_text(
         logger.error(f"{rid} does not exist on DB")
         return {}, {}, {}
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     if extracted_text is None:
-        logger.warning(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.warning(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return {}, {}, {}
 
     split_text: dict[str, str] = extracted_text.split_text
     split_text[MAIN] = extracted_text.text
 
-    split_ners: dict[str, NERS_DICT] = (
-        {}
-    )  # Dict of entity group , with entity and list of positions in field
+    split_ners: dict[
+        str, NERS_DICT
+    ] = {}  # Dict of entity group , with entity and list of positions in field
     split_ners[MAIN] = {}
 
     basic_data = await orm_resource.get_basic()
@@ -138,16 +133,42 @@ async def get_field_text(
                                 split = MAIN
                             else:
                                 split = token.split
-                            split_ners[split].setdefault(token.klass, {}).setdefault(
-                                token.token, []
-                            )
-                            split_ners[split][token.klass][token.token].append(
-                                (token.start, token.end)
-                            )
+                            split_ners[split].setdefault(token.klass, {}).setdefault(token.token, [])
+                            split_ners[split][token.klass][token.token].append((token.start, token.end))
 
     field_metadata = await field_obj.get_field_metadata()
     # Check computed definition of entities
     if field_metadata is not None:
+        # Data Augmentation + Processor entities
+        for data_augmentation_task_id, entities in field_metadata.metadata.entities.items():
+            for entity in entities.entities:
+                entity_text = entity.text
+                entity_label = entity.label
+                entity_positions = entity.positions
+                if entity_label in valid_entity_groups:
+                    split_ners[MAIN].setdefault(entity_label, {}).setdefault(entity_text, [])
+                    for position in entity_positions:
+                        split_ners[MAIN][entity_label][entity_text].append(
+                            (position.start, position.end)
+                        )
+
+        for split, split_metadata in field_metadata.split_metadata.items():
+            for data_augmentation_task_id, entities in split_metadata.entities.items():
+                for entity in entities.entities:
+                    entity_text = entity.text
+                    entity_label = entity.label
+                    entity_positions = entity.positions
+                    if entity_label in valid_entity_groups:
+                        split_ners.setdefault(split, {}).setdefault(entity_label, {}).setdefault(
+                            entity_text, []
+                        )
+                        for position in entity_positions:
+                            split_ners[split][entity_label][entity_text].append(
+                                (position.start, position.end)
+                            )
+
+        # Legacy processor entities
+        # TODO: Remove once processor doesn't use this anymore and remove the positions and ner fields from the message
         for entity_key, positions in field_metadata.metadata.positions.items():
             entities = entity_key.split("/")
             entity_group = entities[0]
@@ -156,9 +177,7 @@ async def get_field_text(
             if entity_group in valid_entity_groups:
                 split_ners[MAIN].setdefault(entity_group, {}).setdefault(entity, [])
                 for position in positions.position:
-                    split_ners[MAIN][entity_group][entity].append(
-                        (position.start, position.end)
-                    )
+                    split_ners[MAIN][entity_group][entity].append((position.start, position.end))
 
         for split, split_metadata in field_metadata.split_metadata.items():
             for entity_key, positions in split_metadata.positions.items():
@@ -166,24 +185,16 @@ async def get_field_text(
                 entity_group = entities[0]
                 entity = "/".join(entities[1:])
                 if entity_group in valid_entity_groups:
-                    split_ners.setdefault(split, {}).setdefault(
-                        entity_group, {}
-                    ).setdefault(entity, [])
+                    split_ners.setdefault(split, {}).setdefault(entity_group, {}).setdefault(entity, [])
                     for position in positions.position:
-                        split_ners[split][entity_group][entity].append(
-                            (position.start, position.end)
-                        )
+                        split_ners[split][entity_group][entity].append((position.start, position.end))
 
     for split, invalid_tokens in invalid_tokens_split.items():
         for token.klass, token.token, token.start, token.end in invalid_tokens:
             if token.klass in split_ners.get(split, {}):
                 if token.token in split_ners.get(split, {}).get(token.klass, {}):
-                    if (token.start, token.end) in split_ners[split][token.klass][
-                        token.token
-                    ]:
-                        split_ners[split][token.klass][token.token].remove(
-                            (token.start, token.end)
-                        )
+                    if (token.start, token.end) in split_ners[split][token.klass][token.token]:
+                        split_ners[split][token.klass][token.token].remove((token.start, token.end))
                         if len(split_ners[split][token.klass][token.token]) == 0:
                             del split_ners[split][token.klass][token.token]
                         if len(split_ners[split][token.klass]) == 0:
@@ -197,9 +208,7 @@ async def get_field_text(
                 for position in positions:
                     split_positions[position] = (entity_group, entity)
 
-        ordered_positions[split] = OrderedDict(
-            sorted(split_positions.items(), key=lambda x: x[0])
-        )
+        ordered_positions[split] = OrderedDict(sorted(split_positions.items(), key=lambda x: x[0]))
 
     split_paragraphs: dict[str, list[tuple[int, int]]] = {}
     if field_metadata is not None:

nucliadb/train/generators/utils.py

@@ -19,19 +19,17 @@
 #
 
 from contextvars import ContextVar
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Optional, Type
 
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
-from nucliadb.ingest.orm.resource import KB_REVERSE
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
-from nucliadb.middleware.transaction import get_read_only_transaction
 from nucliadb.train import SERVICE_NAME, logger
-from nucliadb.train.types import TrainBatchType
+from nucliadb.train.types import T
 from nucliadb_utils.utilities import get_storage
 
-rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar(
-    "rcache", default=None
-)
+rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar("rcache", default=None)
 
 
 def get_resource_cache(clear: bool = False) -> dict[str, ResourceORM]:
@@ -46,12 +44,12 @@ async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> Optional[Resour
     resouce_cache = get_resource_cache()
     orm_resource: Optional[ResourceORM] = None
     if uuid not in resouce_cache:
-        transaction = await get_read_only_transaction()
         storage = await get_storage(service_name=SERVICE_NAME)
-        kb = KnowledgeBoxORM(transaction, storage, kbid)
-        orm_resource = await kb.get(uuid)
-        if orm_resource is not None:
-            resouce_cache[uuid] = orm_resource
+        async with get_driver().transaction(read_only=True) as transaction:
+            kb = KnowledgeBoxORM(transaction, storage, kbid)
+            orm_resource = await kb.get(uuid)
+            if orm_resource is not None:
+                resouce_cache[uuid] = orm_resource
     else:
         orm_resource = resouce_cache.get(uuid)
     return orm_resource
@@ -75,13 +73,11 @@ async def get_paragraph(kbid: str, paragraph_id: str) -> str:
         logger.error(f"{rid} does not exist on DB")
         return ""
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     if extracted_text is None:
-        logger.warning(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.warning(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return ""
 
     if split is not None:
@@ -94,8 +90,8 @@ async def get_paragraph(kbid: str, paragraph_id: str) -> str:
 
 
 async def batchify(
-    producer: AsyncIterator[Any], size: int, batch_klass: TrainBatchType
-):
+    producer: AsyncIterator[Any], size: int, batch_klass: Type[T]
+) -> AsyncGenerator[T, None]:
     # NOTE: we are supposing all protobuffers have a data field
     batch = []
     async for item in producer:

nucliadb/train/lifecycle.py

@@ -18,6 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
 from nucliadb.common.cluster.discovery.utils import (
     setup_cluster_discovery,
     teardown_cluster_discovery,
@@ -33,16 +37,16 @@ from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
 from nucliadb_utils.utilities import start_audit_utility, stop_audit_utility
 
 
-async def initialize() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
-
     await setup_cluster_discovery()
     await start_shard_manager()
     await start_train_grpc(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
 
+    yield
 
-async def finalize() -> None:
     await stop_audit_utility()
     await stop_train_grpc()
     await stop_shard_manager()

nucliadb/train/nodes.py

@@ -19,6 +19,15 @@
 #
 from typing import AsyncIterator, Optional
 
+from nucliadb.common import datamanagers
+from nucliadb.common.cluster import manager
+from nucliadb.common.cluster.base import AbstractIndexNode
+
+# XXX: this keys shouldn't be exposed outside datamanagers
+from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG_BASE
+from nucliadb.common.maindb.driver import Driver, Transaction
+from nucliadb.ingest.orm.entities import EntitiesManager
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb_protos.train_pb2 import (
     GetFieldsRequest,
     GetParagraphsRequest,
@@ -30,15 +39,9 @@ from nucliadb_protos.train_pb2 import (
     TrainSentence,
 )
 from nucliadb_protos.writer_pb2 import ShardObject
-
-from nucliadb.common import datamanagers
-from nucliadb.common.cluster import manager
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.common.maindb.driver import Driver, Transaction
-from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
-from nucliadb.ingest.orm.resource import KB_RESOURCE_SLUG_BASE
+from nucliadb_utils import const
 from nucliadb_utils.storages.storage import Storage
+from nucliadb_utils.utilities import has_feature
 
 
 class TrainShardManager(manager.KBShardManager):
@@ -50,13 +53,13 @@ class TrainShardManager(manager.KBShardManager):
     async def get_reader(self, kbid: str, shard: str) -> tuple[AbstractIndexNode, str]:
         shards = await self.get_shards_by_kbid_inner(kbid)
         try:
-            shard_object: ShardObject = next(
-                filter(lambda x: x.shard == shard, shards.shards)
-            )
+            shard_object: ShardObject = next(filter(lambda x: x.shard == shard, shards.shards))
         except StopIteration:
             raise KeyError("Shard not found")
 
-        node_obj, shard_id = manager.choose_node(shard_object)
+        node_obj, shard_id = manager.choose_node(
+            shard_object, use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid})
+        )
         return node_obj, shard_id
 
     async def get_kb_obj(self, txn: Transaction, kbid: str) -> Optional[KnowledgeBox]:
@@ -69,9 +72,7 @@ class TrainShardManager(manager.KBShardManager):
         kbobj = KnowledgeBox(txn, self.storage, kbid)
         return kbobj
 
-    async def get_kb_entities_manager(
-        self, txn: Transaction, kbid: str
-    ) -> Optional[EntitiesManager]:
+    async def get_kb_entities_manager(self, txn: Transaction, kbid: str) -> Optional[EntitiesManager]:
         kbobj = await self.get_kb_obj(txn, kbid)
         if kbobj is None:
             return None
@@ -79,9 +80,7 @@ class TrainShardManager(manager.KBShardManager):
         manager = EntitiesManager(kbobj, txn)
         return manager
 
-    async def kb_sentences(
-        self, request: GetSentencesRequest
-    ) -> AsyncIterator[TrainSentence]:
+    async def kb_sentences(self, request: GetSentencesRequest) -> AsyncIterator[TrainSentence]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             if request.uuid != "":
@@ -95,24 +94,18 @@ class TrainShardManager(manager.KBShardManager):
                     async for sentence in resource.iterate_sentences(request.metadata):
                         yield sentence
 
-    async def kb_paragraphs(
-        self, request: GetParagraphsRequest
-    ) -> AsyncIterator[TrainParagraph]:
+    async def kb_paragraphs(self, request: GetParagraphsRequest) -> AsyncIterator[TrainParagraph]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             if request.uuid != "":
                 # Filter by uuid
                 resource = await kb.get(request.uuid)
                 if resource:
-                    async for paragraph in resource.iterate_paragraphs(
-                        request.metadata
-                    ):
+                    async for paragraph in resource.iterate_paragraphs(request.metadata):
                         yield paragraph
             else:
                 async for resource in kb.iterate_resources():
-                    async for paragraph in resource.iterate_paragraphs(
-                        request.metadata
-                    ):
+                    async for paragraph in resource.iterate_paragraphs(request.metadata):
                         yield paragraph
 
     async def kb_fields(self, request: GetFieldsRequest) -> AsyncIterator[TrainField]:
@@ -129,15 +122,13 @@ class TrainShardManager(manager.KBShardManager):
                     async for field in resource.iterate_fields(request.metadata):
                         yield field
 
-    async def kb_resources(
-        self, request: GetResourcesRequest
-    ) -> AsyncIterator[TrainResource]:
+    async def kb_resources(self, request: GetResourcesRequest) -> AsyncIterator[TrainResource]:
         async with self.driver.transaction() as txn:
             kb = KnowledgeBox(txn, self.storage, request.kb.uuid)
             base = KB_RESOURCE_SLUG_BASE.format(kbid=request.kb.uuid)
-            async for key in txn.keys(match=base, count=-1):
+            async for key in txn.keys(match=base):
                 # Fetch and Add wanted item
-                rid = await txn.get(key)
+                rid = await txn.get(key, for_update=False)
                 if rid is not None:
                     resource = await kb.get(rid.decode())
                     if resource is not None:

nucliadb/train/servicer.py

@@ -18,10 +18,13 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import traceback
-from typing import Optional
 
 import aiohttp
-from nucliadb_protos.knowledgebox_pb2 import Labels
+
+from nucliadb.common import datamanagers
+from nucliadb.train.settings import settings
+from nucliadb.train.utils import get_shard_manager
+from nucliadb_protos import train_pb2_grpc
 from nucliadb_protos.train_pb2 import (
     GetFieldsRequest,
     GetInfoRequest,
@@ -38,10 +41,6 @@ from nucliadb_protos.writer_pb2 import (
     GetLabelsRequest,
     GetLabelsResponse,
 )
-
-from nucliadb.train.settings import settings
-from nucliadb.train.utils import get_shard_manager
-from nucliadb_protos import train_pb2_grpc
 from nucliadb_telemetry import errors
 
 
@@ -111,20 +110,15 @@ class TrainServicer(train_pb2_grpc.TrainServicer):
     async def GetOntology(  # type: ignore
         self, request: GetLabelsRequest, context=None
     ) -> GetLabelsResponse:
-        async with self.proc.driver.transaction() as txn:
-            kbobj = await self.proc.get_kb_obj(txn, request.kb.uuid)
-            labels: Optional[Labels] = None
-            if kbobj is not None:
-                labels = await kbobj.get_labels()
-
         response = GetLabelsResponse()
-        if kbobj is None:
-            response.status = GetLabelsResponse.Status.NOTFOUND
+        kbid = request.kb.uuid
+        labels = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+        if labels is not None:
+            response.kb.uuid = kbid
+            response.status = GetLabelsResponse.Status.OK
+            response.labels.CopyFrom(labels)
         else:
-            response.kb.uuid = kbobj.kbid
-            if labels is not None:
-                response.labels.CopyFrom(labels)
-
+            response.status = GetLabelsResponse.Status.NOTFOUND
         return response
 
     async def GetOntologyCount(  # type: ignore
@@ -132,9 +126,7 @@ class TrainServicer(train_pb2_grpc.TrainServicer):
     ) -> LabelsetsCount:
         url = settings.internal_search_api.format(kbid=request.kb.uuid)
         facets = [f"faceted=/p/{labelset}" for labelset in request.paragraph_labelsets]
-        facets.extend(
-            [f"faceted=/l/{labelset}" for labelset in request.resource_labelsets]
-        )
+        facets.extend([f"faceted=/l/{labelset}" for labelset in request.resource_labelsets])
         query = "&".join(facets)
         headers = {"X-NUCLIADB-ROLES": "READER"}
         async with aiohttp.ClientSession() as sess:

nucliadb/train/settings.py

@@ -29,13 +29,9 @@ class Settings(DriverSettings):
     nuclia_learning_url: Optional[str] = "https://nuclia.cloud/api/v1/learning/"
     nuclia_learning_apikey: Optional[str] = None
 
-    internal_counter_api: str = (
-        "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/counters"
-    )
+    internal_counter_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/counters"
 
-    internal_search_api: str = (
-        "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
-    )
+    internal_search_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
 
 
 settings = Settings()

nucliadb/train/types.py

@@ -17,7 +17,7 @@
 # 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 Union
+from typing import TypeVar, Union
 
 from nucliadb_protos import dataset_pb2 as dpb
 
@@ -29,14 +29,17 @@ TrainBatch = Union[
     dpb.QuestionAnswerStreamingBatch,
     dpb.SentenceClassificationBatch,
     dpb.TokenClassificationBatch,
+    dpb.FieldStreamingBatch,
 ]
 
-TrainBatchType = Union[
-    type[dpb.FieldClassificationBatch],
-    type[dpb.ImageClassificationBatch],
-    type[dpb.ParagraphClassificationBatch],
-    type[dpb.ParagraphStreamingBatch],
-    type[dpb.QuestionAnswerStreamingBatch],
-    type[dpb.SentenceClassificationBatch],
-    type[dpb.TokenClassificationBatch],
-]
+T = TypeVar(
+    "T",
+    dpb.FieldClassificationBatch,
+    dpb.ImageClassificationBatch,
+    dpb.ParagraphClassificationBatch,
+    dpb.ParagraphStreamingBatch,
+    dpb.QuestionAnswerStreamingBatch,
+    dpb.SentenceClassificationBatch,
+    dpb.TokenClassificationBatch,
+    dpb.FieldStreamingBatch,
+)

nucliadb/train/upload.py

@@ -19,11 +19,10 @@
 #
 import argparse
 import asyncio
+import importlib.metadata
 from asyncio import tasks
 from typing import Callable
 
-import pkg_resources
-
 from nucliadb.train.uploader import start_upload
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
@@ -33,9 +32,7 @@ from nucliadb_utils.settings import running_settings
 def arg_parse():
     parser = argparse.ArgumentParser(description="Upload data to Nuclia Learning API.")
 
-    parser.add_argument(
-        "-r", "--request", dest="request", help="Request UUID", required=True
-    )
+    parser.add_argument("-r", "--request", dest="request", help="Request UUID", required=True)
 
     parser.add_argument("-k", "--kb", dest="kb", help="Knowledge Box", required=True)
 
@@ -75,7 +72,7 @@ def _cancel_all_tasks(loop):
 def run() -> None:
     setup_logging()
 
-    errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
+    errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
     if asyncio._get_running_loop() is not None:
         raise RuntimeError("cannot be called from a running event loop")

nucliadb/train/uploader.py

@@ -20,6 +20,14 @@
 from typing import Optional
 
 import aiohttp
+
+from nucliadb.common import datamanagers
+from nucliadb.common.maindb.utils import setup_driver
+from nucliadb.ingest.orm.entities import EntitiesManager
+from nucliadb.ingest.orm.processor import Processor
+from nucliadb.train import SERVICE_NAME
+from nucliadb.train.models import RequestData
+from nucliadb.train.settings import settings
 from nucliadb_protos.knowledgebox_pb2 import Labels
 from nucliadb_protos.train_pb2 import (
     EnabledMetadata,
@@ -34,13 +42,6 @@ from nucliadb_protos.writer_pb2 import (
     GetLabelsRequest,
     GetLabelsResponse,
 )
-
-from nucliadb.common.maindb.utils import setup_driver
-from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.processor import Processor
-from nucliadb.train import SERVICE_NAME
-from nucliadb.train.models import RequestData
-from nucliadb.train.settings import settings
 from nucliadb_utils.utilities import get_pubsub, get_storage
 
 
@@ -74,9 +75,8 @@ class UploadServicer:
     ) -> GetEntitiesResponse:
         kbid = request.kb.uuid
         response = GetEntitiesResponse()
-        async with self.proc.driver.transaction() as txn:
+        async with self.proc.driver.transaction(read_only=True) as txn:
             kbobj = await self.proc.get_kb_obj(txn, request.kb)
-
             if kbobj is None:
                 response.status = GetEntitiesResponse.Status.NOTFOUND
                 return response
@@ -90,20 +90,16 @@ class UploadServicer:
     async def GetOntology(  # type: ignore
         self, request: GetLabelsRequest, context=None
     ) -> GetLabelsResponse:
-        async with self.proc.driver.transaction() as txn:
-            kbobj = await self.proc.get_kb_obj(txn, request.kb)
-            labels: Optional[Labels] = None
-            if kbobj is not None:
-                labels = await kbobj.get_labels()
-
+        kbid = request.kb.uuid
         response = GetLabelsResponse()
-        if kbobj is None:
+        kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+        if not kb_exists:
             response.status = GetLabelsResponse.Status.NOTFOUND
-        else:
-            response.kb.uuid = kbobj.kbid
-            if labels is not None:
-                response.labels.CopyFrom(labels)
-
+            return response
+        response.kb.uuid = kbid
+        labels: Optional[Labels] = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+        if labels is not None:
+            response.labels.CopyFrom(labels)
         return response
 
 
@@ -123,9 +119,9 @@ async def start_upload(request: str, kb: str):
         }
     ) as sess:
         req = await sess.get(f"{url}/request")
-        request_data = RequestData.parse_raw(await req.read())
+        request_data = RequestData.model_validate_json(await req.read())
 
-        metadata = EnabledMetadata(**request_data.metadata.dict())
+        metadata = EnabledMetadata(**request_data.metadata.model_dump())
 
         if request_data.sentences:
             pbsr = GetSentencesRequest()

nucliadb/train/utils.py

@@ -23,7 +23,7 @@ from grpc import aio
 from grpc_health.v1 import health, health_pb2_grpc
 
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
-from nucliadb.train.nodes import TrainShardManager  # type: ignore
+from nucliadb.train.nodes import TrainShardManager
 from nucliadb.train.settings import settings
 from nucliadb_protos import train_pb2_grpc
 from nucliadb_telemetry.utils import setup_telemetry

nucliadb/writer/__init__.py

@@ -29,9 +29,7 @@ logger = logging.getLogger(SERVICE_NAME)
 class EndpointFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
         return (
-            record.args is not None
-            and len(record.args) >= 3
-            and record.args[2] not in ("/", "/metrics")  # type: ignore
+            record.args is not None and len(record.args) >= 3 and record.args[2] not in ("/", "/metrics")  # type: ignore
         )
 
 

nucliadb/{ingest/fields/keywordset.py → writer/api/utils.py}

@@ -17,18 +17,21 @@
 # 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_protos.resources_pb2 import FieldKeywordset
+from functools import wraps
 
-from nucliadb.ingest.fields.base import Field
+from fastapi import HTTPException
 
+from nucliadb_utils.settings import is_onprem_nucliadb
 
-class Keywordset(Field):
-    pbklass = FieldKeywordset
-    value: FieldKeywordset
-    type: str = "k"
 
-    async def set_value(self, payload: FieldKeywordset):
-        await self.db_set_value(payload)
+def only_for_onprem(fun):
+    @wraps(fun)
+    async def endpoint_wrapper(*args, **kwargs):
+        if not is_onprem_nucliadb():
+            raise HTTPException(
+                status_code=403,
+                detail="This endpoint is only available for onprem NucliaDB",
+            )
+        return await fun(*args, **kwargs)
 
-    async def get_value(self) -> FieldKeywordset:
-        return await self.db_get_value()
+    return endpoint_wrapper