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

nucliadb/common/nidx.py

@@ -0,0 +1,260 @@
+# 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 os
+from typing import Optional
+
+from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxSearcherStub
+
+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
+from nucliadb_utils.settings import FileBackendConfig, indexing_settings, storage_settings
+from nucliadb_utils.storages.settings import settings as extended_storage_settings
+from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
+
+NIDX_ENABLED = bool(os.environ.get("NIDX_ENABLED"))
+
+
+class NidxUtility:
+    api_client = None
+    searcher_client = None
+
+    async def initialize(self):
+        raise NotImplementedError()
+
+    async def finalize(self):
+        raise NotImplementedError()
+
+    async def index(self, msg: IndexMessage) -> int:
+        raise NotImplementedError()
+
+    def wait_for_sync(self):
+        pass
+
+
+def _storage_config(prefix: str, bucket: Optional[str]) -> dict[str, str]:
+    config = {}
+    if storage_settings.file_backend == FileBackendConfig.LOCAL:
+        local_bucket = bucket or storage_settings.local_indexing_bucket
+        file_path = f"{storage_settings.local_files}/{local_bucket}"
+        os.makedirs(file_path, exist_ok=True)
+
+        config[f"{prefix}__OBJECT_STORE"] = "file"
+        config[f"{prefix}__FILE_PATH"] = file_path
+    elif storage_settings.file_backend == FileBackendConfig.GCS:
+        gcs_bucket = bucket or extended_storage_settings.gcs_indexing_bucket
+        config[f"{prefix}__OBJECT_STORE"] = "gcs"
+        if gcs_bucket:
+            config[f"{prefix}__BUCKET"] = gcs_bucket
+        if storage_settings.gcs_base64_creds:
+            config[f"{prefix}__BASE64_CREDS"] = storage_settings.gcs_base64_creds
+        if storage_settings.gcs_endpoint_url:
+            config[f"{prefix}__ENDPOINT"] = storage_settings.gcs_endpoint_url
+    elif storage_settings.file_backend == FileBackendConfig.S3:
+        s3_bucket = bucket or extended_storage_settings.s3_indexing_bucket
+        config[f"{prefix}__OBJECT_STORE"] = "s3"
+        if s3_bucket:
+            config[f"{prefix}__BUCKET"] = s3_bucket
+        config[f"{prefix}__CLIENT_ID"] = storage_settings.s3_client_id or ""
+        config[f"{prefix}__CLIENT_SECRET"] = storage_settings.s3_client_secret or ""
+        config[f"{prefix}__REGION_NAME"] = storage_settings.s3_region_name or ""
+        if storage_settings.s3_endpoint:
+            config[f"{prefix}__ENDPOINT"] = storage_settings.s3_endpoint
+
+    return config
+
+
+class NidxBindingUtility(NidxUtility):
+    """Implements Nidx utility using the binding"""
+
+    def __init__(self):
+        if ingest_settings.driver != DriverConfig.PG:
+            raise ValueError("nidx_binding requires DRIVER=pg")
+
+        self.config = {
+            "METADATA__DATABASE_URL": ingest_settings.driver_pg_url,
+            **_storage_config("INDEXER", None),
+            **_storage_config("STORAGE", "nidx"),
+        }
+
+    async def initialize(self):
+        import nidx_binding  # type: ignore
+
+        self.binding = nidx_binding.NidxBinding(self.config)
+        self.api_client = NidxApiStub(
+            get_traced_grpc_channel(f"localhost:{self.binding.api_port}", "nidx_api")
+        )
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(f"localhost:{self.binding.searcher_port}", "nidx_searcher")
+        )
+
+    async def finalize(self):
+        del self.binding
+
+    async def index(self, msg: IndexMessage) -> int:
+        return self.binding.index(msg.SerializeToString())
+
+    def wait_for_sync(self):
+        self.binding.wait_for_sync()
+
+
+class NidxServiceUtility(NidxUtility):
+    """Implements Nidx utility connecting to the network service"""
+
+    def __init__(self):
+        if indexing_settings.index_nidx_subject is None:
+            raise ValueError("INDEX_NIDX_SUBJECT needed for nidx utility")
+
+        if not settings.nidx_api_address or not settings.nidx_searcher_address:
+            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
+
+        self.nats_connection_manager = NatsConnectionManager(
+            service_name="NidxIndexer",
+            nats_servers=indexing_settings.index_jetstream_servers,
+            nats_creds=indexing_settings.index_jetstream_auth,
+        )
+        self.subject = indexing_settings.index_nidx_subject
+
+    async def initialize(self):
+        await self.nats_connection_manager.initialize()
+        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
+        )
+
+    async def finalize(self):
+        await self.nats_connection_manager.finalize()
+
+    async def index(self, writer: IndexMessage) -> int:
+        res = await self.nats_connection_manager.js.publish(self.subject, writer.SerializeToString())
+        logger.info(
+            f" = Pushed message to nidx shard: {writer.shard}, txid: {writer.txid}  seqid: {res.seq}"  # noqa
+        )
+        return res.seq
+
+
+async def start_nidx_utility() -> Optional[NidxUtility]:
+    if not NIDX_ENABLED:
+        return None
+
+    nidx = get_nidx()
+    if nidx:
+        return nidx
+
+    nidx_utility: NidxUtility
+    if settings.standalone_mode:
+        nidx_utility = NidxBindingUtility()
+    else:
+        nidx_utility = NidxServiceUtility()
+
+    await nidx_utility.initialize()
+    set_utility(Utility.NIDX, nidx_utility)
+    return nidx_utility
+
+
+async def stop_nidx_utility():
+    nidx_utility = get_nidx()
+    if nidx_utility:
+        clean_utility(Utility.NIDX)
+        await nidx_utility.finalize()
+
+
+def get_nidx() -> Optional[NidxUtility]:
+    return get_utility(Utility.NIDX)
+
+
+def get_nidx_api_client() -> Optional["NidxApiStub"]:
+    nidx = get_nidx()
+    if nidx:
+        return nidx.api_client
+    else:
+        return None
+
+
+def get_nidx_searcher_client() -> Optional["NidxSearcherStub"]:
+    nidx = get_nidx()
+    if nidx:
+        return nidx.searcher_client
+    else:
+        return None
+
+
+# TODO: Remove the index node abstraction
+class NodeNidxAdapter:
+    def __init__(self, api_client, searcher_client):
+        # API methods
+        self.GetShard = api_client.GetShard
+        self.NewShard = api_client.NewShard
+        self.DeleteShard = api_client.DeleteShard
+        self.ListShards = api_client.ListShards
+        self.AddVectorSet = api_client.AddVectorSet
+        self.RemoveVectorSet = api_client.RemoveVectorSet
+        self.ListVectorSets = api_client.ListVectorSets
+        self.GetMetadata = api_client.GetMetadata
+
+        # Searcher methods
+        self.Search = searcher_client.Search
+        self.Suggest = searcher_client.Suggest
+        self.Paragraphs = searcher_client.Paragraphs
+        self.Documents = searcher_client.Documents
+
+
+class FakeNode(AbstractIndexNode):
+    def __init__(self, api_client, searcher_client):
+        self.client = NodeNidxAdapter(api_client, searcher_client)
+
+    @property
+    def reader(self):
+        return self.client
+
+    @property
+    def writer(self):
+        return self.client
+
+    def is_read_replica(_):
+        return False
+
+    @property
+    def id(self):
+        return "nidx"
+
+    @property
+    def address(self):
+        return "nidx"
+
+    @property
+    def primary_id(self):
+        return "nidx"
+
+
+def get_nidx_fake_node() -> Optional[FakeNode]:
+    nidx = get_nidx()
+    if nidx:
+        return FakeNode(nidx.api_client, nidx.searcher_client)
+    else:
+        return None

nucliadb/export_import/datamanager.py

@@ -17,6 +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/>.
 #
+import json
 from datetime import datetime
 from typing import AsyncGenerator, Union
 
@@ -53,25 +54,38 @@ class ExportImportDataManager:
 
     async def get_metadata(self, type: str, kbid: str, id: str) -> Metadata:
         key = self._get_maindb_metadata_key(type, kbid, id)
-        async with self.driver.transaction() as txn:
+        async with self.driver.transaction(read_only=True) as txn:
             data = await txn.get(key)
         if data is None or data == b"":
             raise MetadataNotFound()
         decoded = data.decode("utf-8")
-        model_type = {
-            "export": ExportMetadata,
-            "import": ImportMetadata,
-        }[type]
-        return model_type.parse_raw(decoded)  # type: ignore
+        if type == "export":
+            model_type = ExportMetadata
+        elif type == "import":
+            model_type = ImportMetadata  # type: ignore
+        else:
+            raise ValueError(f"Invalid type: {type}")
+        json_decoded = json.loads(decoded)
+
+        # For some reason, the total and processed fields are not always present in the metadata.
+        # This is to unblock already created exports that hit this bug.
+        if json_decoded.get("total") is None:
+            json_decoded["total"] = 0
+        if json_decoded.get("processed") is None:
+            json_decoded["processed"] = 0
+
+        return model_type.model_validate(json_decoded)
 
     async def set_metadata(
         self,
         type: str,
         metadata: Metadata,
     ):
+        metadata.processed = metadata.processed or 0
+        metadata.total = metadata.total or 0
         metadata.modified = datetime.utcnow()
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
-        data = metadata.json().encode("utf-8")
+        data = metadata.model_dump_json().encode("utf-8")
         async with self.driver.transaction() as txn:
             await txn.set(key, data)
             await txn.commit()
@@ -97,9 +111,7 @@ class ExportImportDataManager:
         await self.storage.uploaditerator(export_bytes, field, cf)
         return cf.size
 
-    async def download_export(
-        self, kbid: str, export_id: str
-    ) -> AsyncGenerator[bytes, None]:
+    async def download_export(self, kbid: str, export_id: str) -> AsyncGenerator[bytes, None]:
         key = STORAGE_EXPORT_KEY.format(export_id=export_id)
         bucket = self.storage.get_bucket_name(kbid)
         async for chunk in self.storage.download(bucket, key):
@@ -125,13 +137,9 @@ class ExportImportDataManager:
         async for chunk in self.storage.download(bucket, key):
             yield chunk
 
-    def _get_storage_field(
-        self, kbid: str, key: str, cf: resources_pb2.CloudFile
-    ) -> StorageField:
+    def _get_storage_field(self, kbid: str, key: str, cf: resources_pb2.CloudFile) -> StorageField:
         bucket = self.storage.get_bucket_name(kbid)
-        return self.storage.field_klass(
-            storage=self.storage, bucket=bucket, fullkey=key, field=cf
-        )
+        return self.storage.field_klass(storage=self.storage, bucket=bucket, fullkey=key, field=cf)
 
     async def delete_import(self, kbid: str, import_id: str):
         key = STORAGE_IMPORT_KEY.format(import_id=import_id)
@@ -151,6 +159,4 @@ class ExportImportDataManager:
             await func(kbid, id)
         except Exception as ex:
             errors.capture_exception(ex)
-            logger.exception(
-                f"Could not delete {type} {id} from storage", extra={"kbid": kbid}
-            )
+            logger.exception(f"Could not delete {type} {id} from storage", extra={"kbid": kbid})

nucliadb/export_import/exporter.py

@@ -70,9 +70,7 @@ async def export_kb(
         yield chunk
 
 
-async def export_kb_to_blob_storage(
-    context: ApplicationContext, msg: NatsTaskMessage
-) -> None:
+async def export_kb_to_blob_storage(context: ApplicationContext, msg: NatsTaskMessage) -> None:
     """
     Exports the data of a knowledgebox to the blob storage service.
     """
@@ -90,7 +88,7 @@ async def export_kb_to_blob_storage(
     export_size = await upload_export_retried(iterator, kbid, export_id)
 
     # Store export size
-    metadata.total = metadata.processed = export_size
+    metadata.total = metadata.processed = export_size or 0
     await dm.set_metadata("export", metadata)
 
 
@@ -107,9 +105,7 @@ async def export_resources(
             yield chunk
 
 
-async def export_resources_resumable(
-    context, metadata: ExportMetadata
-) -> AsyncGenerator[bytes, None]:
+async def export_resources_resumable(context, metadata: ExportMetadata) -> AsyncGenerator[bytes, None]:
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
 
     kbid = metadata.kbid
@@ -200,11 +196,9 @@ async def export_learning_config(
 ) -> AsyncGenerator[bytes, None]:
     lconfig = await get_learning_config(kbid)
     if lconfig is None:
-        logger.warning(
-            f"No learning configuration found for kbid", extra={"kbid": kbid}
-        )
+        logger.warning(f"No learning configuration found for kbid", extra={"kbid": kbid})
         return
-    data = lconfig.json().encode("utf-8")
+    data = lconfig.model_dump_json().encode("utf-8")
     yield ExportedItemType.LEARNING_CONFIG.encode("utf-8")
     yield len(data).to_bytes(4, byteorder="big")
     yield data

nucliadb/export_import/importer.py

@@ -89,27 +89,25 @@ async def import_kb(
             await dm.set_metadata("import", metadata)
 
     if metadata is not None:
-        metadata.processed = stream_reader.read_bytes
+        metadata.processed = stream_reader.read_bytes or 0
         await dm.set_metadata("import", metadata)
 
 
-async def import_kb_from_blob_storage(
-    context: ApplicationContext, msg: NatsTaskMessage
-):
+async def import_kb_from_blob_storage(context: ApplicationContext, msg: NatsTaskMessage):
     """
     Imports to a knowledgebox from an export stored in the blob storage service.
     """
     kbid, import_id = msg.kbid, msg.id
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
     metadata = await dm.get_metadata(type="import", kbid=kbid, id=import_id)
-    stream = dm.download_import(kbid, import_id)
 
     retry_handler = TaskRetryHandler("import", dm, metadata)
 
     @retry_handler.wrap
-    async def import_kb_retried(context, kbid, stream, metadata):
+    async def import_kb_retried(context, kbid, metadata):
+        stream = dm.download_import(kbid, import_id)
         await import_kb(context, kbid, stream, metadata)
 
-    await import_kb_retried(context, kbid, stream, metadata)  # type: ignore
+    await import_kb_retried(context, kbid, metadata)
 
     await dm.try_delete_from_storage("import", kbid, import_id)

nucliadb/export_import/models.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 datetime import datetime
+import datetime
 from enum import Enum
 from typing import Any
 
@@ -57,8 +57,8 @@ class Metadata(BaseModel):
     task: TaskMetadata = TaskMetadata(status=Status.SCHEDULED)
     total: int = 0
     processed: int = 0
-    created: datetime = datetime.utcnow()
-    modified: datetime = datetime.utcnow()
+    created: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
+    modified: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
 
 
 class ExportMetadata(Metadata):

nucliadb/export_import/tasks.py

@@ -32,7 +32,7 @@ def get_exports_consumer() -> NatsTaskConsumer:
         name="exports_consumer",
         stream=const.Streams.KB_EXPORTS,  # type: ignore
         callback=export_kb_to_blob_storage,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
@@ -41,7 +41,7 @@ async def get_exports_producer(context: ApplicationContext) -> NatsTaskProducer:
     producer = create_producer(
         name="exports_producer",
         stream=const.Streams.KB_EXPORTS,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)
     return producer
@@ -52,7 +52,7 @@ def get_imports_consumer() -> NatsTaskConsumer:
         name="imports_consumer",
         stream=const.Streams.KB_IMPORTS,  # type: ignore
         callback=import_kb_from_blob_storage,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
@@ -61,7 +61,7 @@ async def get_imports_producer(context: ApplicationContext) -> NatsTaskProducer:
     producer = create_producer(
         name="imports_producer",
         stream=const.Streams.KB_IMPORTS,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)
     return producer

nucliadb/export_import/utils.py

@@ -20,7 +20,6 @@
 import functools
 from typing import AsyncGenerator, AsyncIterator, Callable, Optional
 
-import nats.errors
 from google.protobuf.message import DecodeError as ProtobufDecodeError
 
 from nucliadb import learning_proxy
@@ -34,10 +33,12 @@ from nucliadb.export_import.exceptions import (
     WrongExportStreamFormat,
 )
 from nucliadb.export_import.models import ExportedItemType, ExportItem, Metadata
+from nucliadb.ingest.orm.broker_message import generate_broker_message
 from nucliadb_models.export_import import Status
 from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import resources_pb2, writer_pb2
 from nucliadb_utils.const import Streams
+from nucliadb_utils.transaction import MaxTransactionSizeExceededError
 
 BinaryStream = AsyncGenerator[bytes, None]
 BinaryStreamGenerator = Callable[[int], BinaryStream]
@@ -59,9 +60,6 @@ WRITER_BM_FIELDS = [
     "files",
     "texts",
     "conversations",
-    "layouts",
-    "keywordsets",
-    "datetimes",
 ]
 
 
@@ -88,7 +86,7 @@ async def transaction_commit(
             wait=False,
             target_subject=Streams.INGEST_PROCESSED.subject,
         )
-    except nats.errors.MaxPayloadError:
+    except MaxTransactionSizeExceededError:
         stored_key = await context.blob_storage.set_stream_message(
             kbid=bm.kbid, rid=bm.uuid, data=bm.SerializeToString()
         )
@@ -151,23 +149,17 @@ async def set_entities_groups(
     context: ApplicationContext, kbid: str, entities_groups: kb_pb2.EntitiesGroups
 ) -> None:
     async with datamanagers.with_transaction() as txn:
-        await datamanagers.entities.set_entities_groups(
-            txn, kbid=kbid, entities_groups=entities_groups
-        )
+        await datamanagers.entities.set_entities_groups(txn, kbid=kbid, entities_groups=entities_groups)
         await txn.commit()
 
 
-async def set_labels(
-    context: ApplicationContext, kbid: str, labels: kb_pb2.Labels
-) -> None:
+async def set_labels(context: ApplicationContext, kbid: str, labels: kb_pb2.Labels) -> None:
     async with datamanagers.with_transaction() as txn:
         await datamanagers.labels.set_labels(txn, kbid=kbid, labels=labels)
         await txn.commit()
 
 
-async def iter_kb_resource_uuids(
-    context: ApplicationContext, kbid: str
-) -> AsyncGenerator[str, None]:
+async def iter_kb_resource_uuids(context: ApplicationContext, kbid: str) -> AsyncGenerator[str, None]:
     async for rid in datamanagers.resources.iterate_resource_ids(kbid=kbid):
         yield rid
 
@@ -175,8 +167,13 @@ async def iter_kb_resource_uuids(
 async def get_broker_message(
     context: ApplicationContext, kbid: str, rid: str
 ) -> Optional[writer_pb2.BrokerMessage]:
-    async with datamanagers.with_transaction() as txn:
-        return await datamanagers.resources.get_broker_message(txn, kbid=kbid, rid=rid)
+    async with datamanagers.with_ro_transaction() as txn:
+        resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=rid)
+        if resource is None:
+            return None
+        resource.disable_vectors = False
+        resource.txn = txn
+        return await generate_broker_message(resource)
 
 
 def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFile]:
@@ -184,6 +181,10 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
     binaries: list[resources_pb2.CloudFile] = []
     for file_field in bm.files.values():
         if file_field.HasField("file"):
+            if file_field.file.source == resources_pb2.CloudFile.Source.EXTERNAL:
+                # Binaries of externally hosted files are not
+                # to be downloaded and stored in the export file
+                continue
             _clone_collect_cf(binaries, file_field.file)
 
     for conversation in bm.conversations.values():
@@ -191,11 +192,6 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
             for attachment in message.content.attachments:
                 _clone_collect_cf(binaries, attachment)
 
-    for layout in bm.layouts.values():
-        for block in layout.body.blocks.values():
-            if block.HasField("file"):
-                _clone_collect_cf(binaries, block.file)
-
     for field_extracted_data in bm.file_extracted_data:
         if field_extracted_data.HasField("file_thumbnail"):
             _clone_collect_cf(binaries, field_extracted_data.file_thumbnail)
@@ -227,9 +223,7 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
     return binaries
 
 
-def _clone_collect_cf(
-    binaries: list[resources_pb2.CloudFile], origin: resources_pb2.CloudFile
-):
+def _clone_collect_cf(binaries: list[resources_pb2.CloudFile], origin: resources_pb2.CloudFile):
     cf = resources_pb2.CloudFile()
     cf.CopyFrom(origin)
     # Mark the cloud file of the broker message being exported as export source
@@ -246,12 +240,12 @@ async def download_binary(
 
 
 async def get_entities(context: ApplicationContext, kbid: str) -> kb_pb2.EntitiesGroups:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.entities.get_entities_groups(txn, kbid=kbid)
 
 
 async def get_labels(context: ApplicationContext, kbid: str) -> kb_pb2.Labels:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.labels.get_labels(txn, kbid=kbid)
 
 
@@ -398,7 +392,7 @@ class ExportStreamReader:
             # Backward compatible code for old exports that don't have a learning config.
             return None, type_bytes + self.stream.buffer
         data = await self.read_item()
-        lconfig = learning_proxy.LearningConfiguration.parse_raw(data)
+        lconfig = learning_proxy.LearningConfiguration.model_validate_json(data)
         return lconfig, self.stream.buffer
 
     async def iter_items(self) -> AsyncGenerator[ExportItem, None]:
@@ -411,7 +405,7 @@ class ExportStreamReader:
                     ExportedItemType.ENTITIES: self.read_entities,
                     ExportedItemType.LABELS: self.read_labels,
                 }[item_type]
-                data = await read_data_func()  # type: ignore
+                data = await read_data_func()
                 yield item_type, data
             except ExportStreamExhausted:
                 break
@@ -506,21 +500,15 @@ def stream_compatible_with_kb(
     return wrapped()
 
 
-async def _check_semantic_model_compatibility(
-    kbid: str, stream: AsyncGenerator[bytes, None]
-) -> bytes:
+async def _check_semantic_model_compatibility(kbid: str, stream: AsyncGenerator[bytes, None]) -> bytes:
     stream_reader = ExportStreamReader(stream)
     lconfig, leftover_bytes = await stream_reader.maybe_read_learning_config()
     if lconfig is None:
-        logger.warning(
-            "Learning config not found on the export stream. Export may be incompatible."
-        )
+        logger.warning("Learning config not found on the export stream. Export may be incompatible.")
         return leftover_bytes
     kb_lconfig = await get_learning_config(kbid)
     if kb_lconfig is None:
-        logger.warning(
-            "No learning config found on the knowledge box. Export may be incompatible."
-        )
+        logger.warning("No learning config found on the knowledge box. Export may be incompatible.")
         return leftover_bytes
     if kb_lconfig.semantic_model == lconfig.semantic_model:
         logger.info(f"Semantic model match: {kb_lconfig.semantic_model}")

nucliadb/health.py

@@ -78,9 +78,7 @@ async def grpc_health_check(health_servicer) -> None:
         for check in _health_checks:
             if not check():
                 logger.info(f"Health check failed on {check.__name__}")
-                await health_servicer.set(
-                    "", health_pb2.HealthCheckResponse.NOT_SERVING
-                )
+                await health_servicer.set("", health_pb2.HealthCheckResponse.NOT_SERVING)
                 break
         else:
             await health_servicer.set("", health_pb2.HealthCheckResponse.SERVING)