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

nucliadb/train/generators/image_classifier.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 from nucliadb.train.generators.utils import batchify
 from nucliadb_models.filters import FilterExpression
@@ -33,7 +33,7 @@ def image_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[ImageClassificationBatch, None]:
     generator = generate_image_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ImageClassificationBatch)

nucliadb/train/generators/paragraph_classifier.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 from fastapi import HTTPException
 from nidx_protos.nodereader_pb2 import StreamRequest
@@ -38,7 +38,7 @@ def paragraph_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[ParagraphClassificationBatch, None]:
     if len(trainset.filter.labels) != 1:
         raise HTTPException(

nucliadb/train/generators/paragraph_streaming.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
@@ -38,7 +38,7 @@ def paragraph_streaming_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[ParagraphStreamingBatch, None]:
     generator = generate_paragraph_streaming_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ParagraphStreamingBatch)

nucliadb/train/generators/question_answer_streaming.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
@@ -47,7 +47,7 @@ def question_answer_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[QuestionAnswerStreamingBatch, None]:
     generator = generate_question_answer_streaming_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, QuestionAnswerStreamingBatch)

nucliadb/train/generators/sentence_classifier.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 from fastapi import HTTPException
 from nidx_protos.nodereader_pb2 import StreamRequest
@@ -40,7 +40,7 @@ def sentence_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[SentenceClassificationBatch, None]:
     if len(trainset.filter.labels) == 0:
         raise HTTPException(
@@ -116,10 +116,7 @@ async def get_sentences(kbid: str, result: str) -> list[str]:
     if split is not None:
         text = extracted_text.split_text[split]
         for paragraph in field_metadata.split_metadata[split].paragraphs:
-            if paragraph.key == "":
-                key = f"{rid}/{field_type}/{field}/{paragraph.start}-{paragraph.end}"
-            else:
-                key = paragraph.key
+            key = f"{rid}/{field_type}/{field}/{paragraph.start}-{paragraph.end}"
             if key == result:
                 for sentence in paragraph.sentences:
                     splitted_text = text[sentence.start : sentence.end]
@@ -127,10 +124,7 @@ async def get_sentences(kbid: str, result: str) -> list[str]:
     else:
         text = extracted_text.text
         for paragraph in field_metadata.metadata.paragraphs:
-            if paragraph.key == "":
-                key = f"{rid}/{field_type}/{field}/{paragraph.start}-{paragraph.end}"
-            else:
-                key = paragraph.key
+            key = f"{rid}/{field_type}/{field}/{paragraph.start}-{paragraph.end}"
             if key == result:
                 for sentence in paragraph.sentences:
                     splitted_text = text[sentence.start : sentence.end]

nucliadb/train/generators/token_classifier.py

@@ -19,7 +19,8 @@
 #
 
 from collections import OrderedDict
-from typing import AsyncGenerator, Optional, cast
+from collections.abc import AsyncGenerator
+from typing import cast
 
 from nidx_protos.nodereader_pb2 import StreamFilter, StreamRequest
 
@@ -43,7 +44,7 @@ def token_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[TokenClassificationBatch, None]:
     generator = generate_token_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, TokenClassificationBatch)

nucliadb/train/generators/utils.py

@@ -18,7 +18,8 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Any, AsyncGenerator, AsyncIterator, Optional, Type
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any
 
 from nucliadb.common.cache import get_resource_cache
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
@@ -30,16 +31,16 @@ from nucliadb.train.types import T
 from nucliadb_utils.utilities import get_storage
 
 
-async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> Optional[ResourceORM]:
+async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> ResourceORM | None:
     resource_cache = get_resource_cache()
     if resource_cache is None:
-        return await _get_resource_from_db(kbid, uuid)
         logger.warning("Resource cache is not set")
+        return await _get_resource_from_db(kbid, uuid)
 
     return await resource_cache.get(kbid, uuid)
 
 
-async def _get_resource_from_db(kbid: str, uuid: str) -> Optional[ResourceORM]:
+async def _get_resource_from_db(kbid: str, uuid: str) -> ResourceORM | None:
     storage = await get_storage(service_name=SERVICE_NAME)
     async with get_driver().ro_transaction() as transaction:
         kb = KnowledgeBoxORM(transaction, storage, kbid)
@@ -81,7 +82,7 @@ async def get_paragraph(kbid: str, paragraph_id: str) -> str:
 
 
 async def batchify(
-    producer: AsyncIterator[Any], size: int, batch_klass: Type[T]
+    producer: AsyncIterator[Any], size: int, batch_klass: type[T]
 ) -> AsyncGenerator[T, None]:
     # NOTE: we are supposing all protobuffers have a data field
     batch = []

nucliadb/train/nodes.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 AsyncIterator, Optional
+from collections.abc import AsyncIterator
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster import manager
@@ -62,7 +62,7 @@ class TrainShardManager(manager.KBShardManager):
 
         return shard_object.nidx_shard_id
 
-    async def get_kb_obj(self, txn: Transaction, kbid: str) -> Optional[KnowledgeBox]:
+    async def get_kb_obj(self, txn: Transaction, kbid: str) -> KnowledgeBox | None:
         if kbid is None:
             return None
 
@@ -72,7 +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) -> EntitiesManager | None:
         kbobj = await self.get_kb_obj(txn, kbid)
         if kbobj is None:
             return None

nucliadb/train/resource.py

@@ -19,7 +19,7 @@
 #
 from __future__ import annotations
 
-from typing import AsyncIterator, MutableMapping, Optional
+from collections.abc import AsyncIterator, MutableMapping
 
 from nucliadb.common import datamanagers
 from nucliadb.ingest.orm.resource import Resource
@@ -69,9 +69,7 @@ async def iterate_sentences(
             # return any
             vectorset_id = None
             async with datamanagers.with_ro_transaction() as txn:
-                async for vectorset_id, vs in datamanagers.vectorsets.iter(
-                    txn=txn, kbid=resource.kb.kbid
-                ):
+                async for vectorset_id, vs in datamanagers.vectorsets.iter(txn=txn, kbid=resource.kbid):
                     break
             assert vectorset_id is not None, "All KBs must have at least a vectorset"
             vo = await field.get_vectors(vectorset_id, vs.storage_key_kind)
@@ -81,7 +79,7 @@ async def iterate_sentences(
         if fm is None:
             continue
 
-        field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
+        field_metadatas: list[tuple[str | None, FieldMetadata]] = [(None, fm.metadata)]
         for subfield_metadata, splitted_metadata in fm.split_metadata.items():
             field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -188,7 +186,7 @@ async def iterate_paragraphs(
         if fm is None:
             continue
 
-        field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
+        field_metadatas: list[tuple[str | None, FieldMetadata]] = [(None, fm.metadata)]
         for subfield_metadata, splitted_metadata in fm.split_metadata.items():
             field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -264,7 +262,7 @@ async def iterate_fields(
         if fm is None:
             continue
 
-        field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
+        field_metadatas: list[tuple[str | None, FieldMetadata]] = [(None, fm.metadata)]
         for subfield_metadata, splitted_metadata in fm.split_metadata.items():
             field_metadatas.append((subfield_metadata, splitted_metadata))
 
@@ -319,7 +317,7 @@ async def generate_train_resource(
         if fm is None:
             continue
 
-        field_metadatas: list[tuple[Optional[str], FieldMetadata]] = [(None, fm.metadata)]
+        field_metadatas: list[tuple[str | None, FieldMetadata]] = [(None, fm.metadata)]
         for subfield_metadata, splitted_metadata in fm.split_metadata.items():
             field_metadatas.append((subfield_metadata, splitted_metadata))
 

nucliadb/train/settings.py

@@ -17,17 +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 typing import Optional
 
 from nucliadb.ingest.settings import DriverSettings
 
 
 class Settings(DriverSettings):
     grpc_port: int = 8031
-    train_grpc_address: Optional[str] = None
+    train_grpc_address: str | None = None
 
-    nuclia_learning_url: Optional[str] = "https://nuclia.cloud/api/v1/learning/"
-    nuclia_learning_apikey: Optional[str] = None
+    nuclia_learning_url: str | None = "https://nuclia.cloud/api/v1/learning/"
+    nuclia_learning_apikey: str | None = None
 
     internal_counter_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/counters"
 

nucliadb/train/types.py

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

nucliadb/train/upload.py

@@ -20,8 +20,9 @@
 import argparse
 import asyncio
 import importlib.metadata
+import inspect
 from asyncio import tasks
-from typing import Callable
+from collections.abc import Callable
 
 from nucliadb.train.uploader import start_upload
 from nucliadb_telemetry import errors
@@ -89,7 +90,7 @@ def run() -> None:
     finally:
         try:
             for finalizer in finalizers:
-                if asyncio.iscoroutinefunction(finalizer):
+                if inspect.iscoroutinefunction(finalizer):
                     loop.run_until_complete(finalizer())
                 else:
                     finalizer()

nucliadb/train/uploader.py

@@ -17,7 +17,6 @@
 # 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 aiohttp
 
@@ -97,7 +96,7 @@ class UploadServicer:
             response.status = GetLabelsResponse.Status.NOTFOUND
             return response
         response.kb.uuid = kbid
-        labels: Optional[Labels] = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+        labels: Labels | None = await datamanagers.atomic.labelset.get_all(kbid=kbid)
         if labels is not None:
             response.labels.CopyFrom(labels)
         return response

nucliadb/train/utils.py

@@ -17,7 +17,6 @@
 # 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
 
 from grpc import aio
 from grpc_health.v1 import health, health_pb2_grpc
@@ -37,7 +36,7 @@ from nucliadb_utils.utilities import (
 )
 
 
-async def start_train_grpc(service_name: Optional[str] = None):
+async def start_train_grpc(service_name: str | None = None):
     actual_service = get_utility(Utility.TRAIN)
     if actual_service is not None:
         return

nucliadb/writer/api/v1/export_import.py

@@ -17,8 +17,8 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+from collections.abc import AsyncGenerator
 from datetime import datetime
-from typing import AsyncGenerator
 from uuid import uuid4
 
 from fastapi_versioning import version
@@ -66,6 +66,7 @@ from nucliadb_utils.authentication import requires_one
     summary="Start an export of a Knowledge Box",
     tags=["Knowledge Boxes"],
     response_model=CreateExportResponse,
+    include_in_schema=False,
 )
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
 @version(1)
@@ -91,6 +92,7 @@ async def start_kb_export_endpoint(request: Request, kbid: str):
     tags=["Knowledge Boxes"],
     response_model=NewImportedKbResponse,
     openapi_extra={"x-hidden-operation": True},
+    include_in_schema=False,
 )
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
 @version(1)
@@ -140,6 +142,7 @@ async def kb_create_and_import_endpoint(request: Request):
     summary="Start an import to a Knowledge Box",
     tags=["Knowledge Boxes"],
     response_model=CreateImportResponse,
+    include_in_schema=False,
 )
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
 @version(1)

nucliadb/writer/api/v1/field.py

@@ -17,8 +17,9 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+from collections.abc import Callable
 from inspect import iscoroutinefunction
-from typing import TYPE_CHECKING, Annotated, Callable, Optional, Type, Union
+from typing import TYPE_CHECKING, Annotated
 
 import pydantic
 from fastapi import HTTPException, Query, Response
@@ -72,12 +73,7 @@ if TYPE_CHECKING:  # pragma: no cover
 else:
     FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP: dict[models.FieldTypeName, int]
 
-FieldModelType = Union[
-    models.TextField,
-    models.LinkField,
-    models.InputConversationField,
-    models.FileField,
-]
+FieldModelType = models.TextField | models.LinkField | models.InputConversationField | models.FileField
 
 FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP = {
     models.FieldTypeName.FILE: resources_pb2.FieldType.FILE,
@@ -249,9 +245,10 @@ async def parse_conversation_field_adapter(
     writer: BrokerMessage,
     toprocess: PushPayload,
     resource_classifications: ResourceClassifications,
+    replace_field: bool = False,
 ):
     return await parse_conversation_field(
-        field_id, field_payload, writer, toprocess, kbid, rid, resource_classifications
+        field_id, field_payload, writer, toprocess, kbid, rid, resource_classifications, replace_field
     )
 
 
@@ -277,7 +274,7 @@ async def parse_file_field_adapter(
     )
 
 
-FIELD_PARSERS_MAP: dict[Type, Callable] = {
+FIELD_PARSERS_MAP: dict[type, Callable] = {
     models.TextField: parse_text_field_adapter,
     models.LinkField: parse_link_field_adapter,
     models.InputConversationField: parse_conversation_field_adapter,
@@ -380,7 +377,9 @@ async def add_resource_field_conversation_rslug_prefix(
     field_id: FieldIdString,
     field_payload: models.InputConversationField,
 ) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field_payload)
+    return await add_field_to_resource_by_slug(
+        request, kbid, rslug, field_id, field_payload, replace_field=True
+    )
 
 
 @api.put(
@@ -399,7 +398,7 @@ async def add_resource_field_conversation_rid_prefix(
     field_id: FieldIdString,
     field_payload: models.InputConversationField,
 ) -> ResourceFieldAdded:
-    return await add_field_to_resource(request, kbid, rid, field_id, field_payload)
+    return await add_field_to_resource(request, kbid, rid, field_id, field_payload, replace_field=True)
 
 
 @api.put(
@@ -466,7 +465,9 @@ async def append_messages_to_conversation_field_rslug_prefix(
         field = models.InputConversationField(messages=messages)
     except pydantic.ValidationError as e:
         raise HTTPException(status_code=422, detail=str(e))
-    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field)
+    return await add_field_to_resource_by_slug(
+        request, kbid, rslug, field_id, field, replace_field=False
+    )
 
 
 @api.put(
@@ -489,7 +490,7 @@ async def append_messages_to_conversation_field_rid_prefix(
         field = models.InputConversationField(messages=messages)
     except pydantic.ValidationError as e:
         raise HTTPException(status_code=422, detail=str(e))
-    return await add_field_to_resource(request, kbid, rid, field_id, field)
+    return await add_field_to_resource(request, kbid, rid, field_id, field, replace_field=False)
 
 
 @api.delete(
@@ -545,7 +546,7 @@ async def reprocess_file_field(
     rid: str,
     field_id: FieldIdString,
     x_nucliadb_user: Annotated[str, X_NUCLIADB_USER] = "",
-    x_file_password: Annotated[Optional[str], X_FILE_PASSWORD] = None,
+    x_file_password: Annotated[str | None, X_FILE_PASSWORD] = None,
     reset_title: bool = Query(
         default=False,
         description="Reset the title of the resource so that the file or link computed titles are set after processing.",

nucliadb/writer/api/v1/knowledgebox.py

@@ -17,10 +17,9 @@
 # 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 asyncio
 from functools import partial
 
-from fastapi import HTTPException
+from fastapi import BackgroundTasks, HTTPException
 from fastapi_versioning import version
 from starlette.requests import Request
 
@@ -36,10 +35,6 @@ from nucliadb.writer import logger
 from nucliadb.writer.api.utils import only_for_onprem
 from nucliadb.writer.api.v1.router import KB_PREFIX, KBS_PREFIX, api
 from nucliadb.writer.utilities import get_processing
-from nucliadb_models.external_index_providers import (
-    ExternalIndexProviderType,
-    PineconeServerlessCloud,
-)
 from nucliadb_models.resource import (
     KnowledgeBoxConfig,
     KnowledgeBoxObj,
@@ -118,20 +113,6 @@ async def create_kb(item: KnowledgeBoxConfig) -> tuple[str, str]:
     external_index_provider = knowledgebox_pb2.CreateExternalIndexProviderMetadata(
         type=knowledgebox_pb2.ExternalIndexProviderType.UNSET,
     )
-    if (
-        item.external_index_provider
-        and item.external_index_provider.type == ExternalIndexProviderType.PINECONE
-    ):
-        pinecone_api_key = item.external_index_provider.api_key
-        serverless_pb = to_pinecone_serverless_cloud_pb(item.external_index_provider.serverless_cloud)
-        external_index_provider = knowledgebox_pb2.CreateExternalIndexProviderMetadata(
-            type=knowledgebox_pb2.ExternalIndexProviderType.PINECONE,
-            pinecone_config=knowledgebox_pb2.CreatePineconeConfig(
-                api_key=pinecone_api_key,
-                serverless_cloud=serverless_pb,
-            ),
-        )
-
     try:
         (kbid, slug) = await KnowledgeBox.create(
             driver,
@@ -165,8 +146,6 @@ async def create_kb(item: KnowledgeBoxConfig) -> tuple[str, str]:
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
 async def update_kb(request: Request, kbid: str, item: KnowledgeBoxConfig) -> KnowledgeBoxObjID:
-    driver = get_driver()
-    config = None
     if (
         item.slug
         or item.title
@@ -174,29 +153,24 @@ async def update_kb(request: Request, kbid: str, item: KnowledgeBoxConfig) -> Kn
         or item.hidden_resources_enabled
         or item.hidden_resources_hide_on_creation
     ):
-        config = knowledgebox_pb2.KnowledgeBoxConfig(
-            slug=item.slug or "",
-            title=item.title or "",
-            description=item.description or "",
-            hidden_resources_enabled=item.hidden_resources_enabled,
-            hidden_resources_hide_on_creation=item.hidden_resources_hide_on_creation,
-        )
-    try:
-        async with driver.rw_transaction() as txn:
+        try:
+            driver = get_driver()
             await KnowledgeBox.update(
-                txn,
-                uuid=kbid,
+                driver,
+                kbid=kbid,
                 slug=item.slug,
-                config=config,
+                title=item.title,
+                description=item.description,
+                hidden_resources_enabled=item.hidden_resources_enabled,
+                hidden_resources_hide_on_creation=item.hidden_resources_hide_on_creation,
             )
-            await txn.commit()
-    except datamanagers.exceptions.KnowledgeBoxNotFound:
-        raise HTTPException(status_code=404, detail="Knowledge box does not exist")
-    except Exception as exc:
-        logger.exception("Could not update KB", exc_info=exc, extra={"kbid": kbid})
-        raise HTTPException(status_code=500, detail="Error updating knowledge box")
-    else:
-        return KnowledgeBoxObjID(uuid=kbid)
+        except datamanagers.exceptions.KnowledgeBoxNotFound:
+            raise HTTPException(status_code=404, detail="Knowledge box does not exist")
+        except Exception as exc:
+            logger.exception("Could not update KB", exc_info=exc, extra={"kbid": kbid})
+            raise HTTPException(status_code=500, detail="Error updating knowledge box")
+
+    return KnowledgeBoxObjID(uuid=kbid)
 
 
 @only_for_onprem
@@ -209,7 +183,7 @@ async def update_kb(request: Request, kbid: str, item: KnowledgeBoxConfig) -> Kn
 )
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
-async def delete_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
+async def delete_kb(request: Request, kbid: str, background: BackgroundTasks) -> KnowledgeBoxObj:
     driver = get_driver()
     try:
         await KnowledgeBox.delete(driver, kbid=kbid)
@@ -233,18 +207,6 @@ async def delete_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
     # be nice and notify processing this KB is being deleted so we waste
     # resources
     processing = get_processing()
-    asyncio.create_task(processing.delete_from_processing(kbid=kbid))
+    background.add_task(processing.delete_from_processing, kbid=kbid)
 
     return KnowledgeBoxObj(uuid=kbid)
-
-
-def to_pinecone_serverless_cloud_pb(
-    serverless: PineconeServerlessCloud,
-) -> knowledgebox_pb2.PineconeServerlessCloud.ValueType:
-    return {
-        PineconeServerlessCloud.AWS_EU_WEST_1: knowledgebox_pb2.PineconeServerlessCloud.AWS_EU_WEST_1,
-        PineconeServerlessCloud.AWS_US_EAST_1: knowledgebox_pb2.PineconeServerlessCloud.AWS_US_EAST_1,
-        PineconeServerlessCloud.AWS_US_WEST_2: knowledgebox_pb2.PineconeServerlessCloud.AWS_US_WEST_2,
-        PineconeServerlessCloud.AZURE_EASTUS2: knowledgebox_pb2.PineconeServerlessCloud.AZURE_EASTUS2,
-        PineconeServerlessCloud.GCP_US_CENTRAL1: knowledgebox_pb2.PineconeServerlessCloud.GCP_US_CENTRAL1,
-    }[serverless]

nucliadb/writer/api/v1/learning_config.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 fastapi import Request
+from fastapi import Header, Request
 from fastapi_versioning import version
 from nuclia_models.config.proto import ExtractConfig, SplitConfiguration
 
@@ -55,10 +55,11 @@ async def set_configuration(
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
 @version(1)
 async def patch_configuration(
-    request: Request,
-    kbid: str,
+    request: Request, kbid: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
 ):
-    return await learning_config_proxy(request, "PATCH", f"/config/{kbid}")
+    return await learning_config_proxy(
+        request, "PATCH", f"/config/{kbid}", headers={"account-id": x_nucliadb_account}
+    )
 
 
 @api.post(