nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

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, List, 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,
@@ -278,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,
@@ -463,7 +459,7 @@ async def append_messages_to_conversation_field_rslug_prefix(
     kbid: str,
     rslug: str,
     field_id: FieldIdString,
-    messages: List[models.InputMessage],
+    messages: list[models.InputMessage],
 ) -> ResourceFieldAdded:
     try:
         field = models.InputConversationField(messages=messages)
@@ -488,7 +484,7 @@ async def append_messages_to_conversation_field_rid_prefix(
     kbid: str,
     rid: str,
     field_id: FieldIdString,
-    messages: List[models.InputMessage],
+    messages: list[models.InputMessage],
 ) -> ResourceFieldAdded:
     try:
         field = models.InputConversationField(messages=messages)
@@ -550,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
 
@@ -184,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)
@@ -208,6 +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)

nucliadb/writer/api/v1/resource.py

@@ -17,13 +17,12 @@
 # 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
 import contextlib
 from time import time
-from typing import Annotated, Optional
+from typing import Annotated
 from uuid import uuid4
 
-from fastapi import HTTPException, Query, Response
+from fastapi import BackgroundTasks, HTTPException, Query, Response
 from fastapi_versioning import version
 from starlette.requests import Request
 
@@ -498,12 +497,10 @@ async def _reprocess_resource(
 @requires(NucliaDBRoles.WRITER)
 @version(1)
 async def delete_resource_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
+    request: Request, kbid: str, rslug: str, background: BackgroundTasks
 ):
     rid = await get_rid_from_slug_or_raise_error(kbid, rslug)
-    return await _delete_resource(request, kbid, rid)
+    return await _delete_resource(request, kbid, rid, background)
 
 
 @api.delete(
@@ -514,19 +511,11 @@ async def delete_resource_rslug_prefix(
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
-async def delete_resource_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-):
-    return await _delete_resource(request, kbid, rid)
+async def delete_resource_rid_prefix(request: Request, kbid: str, rid: str, background: BackgroundTasks):
+    return await _delete_resource(request, kbid, rid, background)
 
 
-async def _delete_resource(
-    request: Request,
-    kbid: str,
-    rid: str,
-):
+async def _delete_resource(request: Request, kbid: str, rid: str, background: BackgroundTasks):
     await validate_rid_exists_or_raise_error(kbid, rid)
 
     partitioning = get_partitioning()
@@ -541,7 +530,7 @@ async def _delete_resource(
     parse_audit(writer.audit, request)
     await transaction.commit(writer, partition)
     processing = get_processing()
-    asyncio.create_task(processing.delete_from_processing(kbid=kbid, resource_id=rid))
+    background.add_task(processing.delete_from_processing, kbid=kbid, resource_id=rid)
 
     return Response(status_code=204)
 
@@ -637,7 +626,7 @@ def needs_resource_reindex(item: UpdateResourcePayload) -> bool:
     )
 
 
-async def maybe_send_to_process(toprocess: PushPayload, partition) -> Optional[int]:
+async def maybe_send_to_process(toprocess: PushPayload, partition) -> int | None:
     if not needs_reprocess(toprocess):
         return None
 

nucliadb/writer/api/v1/services.py

@@ -17,152 +17,22 @@
 # 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 HTTPException, Response
+from fastapi import Body, HTTPException, Path, Response
 from fastapi_versioning import version
 from starlette.requests import Request
 
 from nucliadb.common import datamanagers
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.models_utils import to_proto
-from nucliadb.models.responses import (
-    HTTPConflict,
-    HTTPInternalServerError,
-    HTTPNotFound,
-)
 from nucliadb.writer.api.v1.router import KB_PREFIX, api
 from nucliadb_models.configuration import SearchConfiguration
-from nucliadb_models.entities import (
-    CreateEntitiesGroupPayload,
-    UpdateEntitiesGroupPayload,
-)
 from nucliadb_models.labels import LabelSet
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.synonyms import KnowledgeBoxSynonyms
 from nucliadb_protos import writer_pb2
 from nucliadb_protos.knowledgebox_pb2 import Label as LabelPB
 from nucliadb_protos.knowledgebox_pb2 import LabelSet as LabelSetPB
-from nucliadb_protos.writer_pb2 import (
-    DelEntitiesRequest,
-    NewEntitiesGroupRequest,
-    NewEntitiesGroupResponse,
-    OpStatusWriter,
-    UpdateEntitiesGroupRequest,
-    UpdateEntitiesGroupResponse,
-)
 from nucliadb_utils.authentication import requires
-from nucliadb_utils.utilities import get_ingest
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/entitiesgroups",
-    status_code=200,
-    summary="Create Knowledge Box Entities Group",
-    tags=["Knowledge Box Services"],
-    openapi_extra={"x-operation_order": 1},
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def create_entities_group(request: Request, kbid: str, item: CreateEntitiesGroupPayload):
-    ingest = get_ingest()
-
-    pbrequest: NewEntitiesGroupRequest = NewEntitiesGroupRequest()
-    pbrequest.kb.uuid = kbid
-    pbrequest.group = item.group
-    pbrequest.entities.custom = True
-    if item.title:
-        pbrequest.entities.title = item.title
-    if item.color:
-        pbrequest.entities.color = item.color
-
-    for key, entity in item.entities.items():
-        entitypb = pbrequest.entities.entities[key]
-        entitypb.value = entity.value
-        entitypb.merged = entity.merged
-        entitypb.deleted = False
-        entitypb.represents.extend(entity.represents)
-
-    status: NewEntitiesGroupResponse = await ingest.NewEntitiesGroup(pbrequest)  # type: ignore
-    if status.status == NewEntitiesGroupResponse.Status.OK:
-        return
-    elif status.status == NewEntitiesGroupResponse.Status.KB_NOT_FOUND:
-        return HTTPNotFound(detail="Knowledge Box does not exist")
-    elif status.status == NewEntitiesGroupResponse.Status.ALREADY_EXISTS:
-        return HTTPConflict(
-            detail=f"Entities group {item.group} already exists in this Knowledge box",
-        )
-    elif status.status == NewEntitiesGroupResponse.Status.ERROR:
-        return HTTPInternalServerError(detail="Error on settings entities on a Knowledge box")
-
-
-@api.patch(
-    f"/{KB_PREFIX}/{{kbid}}/entitiesgroup/{{group}}",
-    status_code=200,
-    summary="Update Knowledge Box Entities Group",
-    tags=["Knowledge Box Services"],
-    openapi_extra={"x-operation_order": 2},
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def update_entities_group(
-    request: Request, kbid: str, group: str, item: UpdateEntitiesGroupPayload
-):
-    ingest = get_ingest()
-
-    pbrequest: UpdateEntitiesGroupRequest = UpdateEntitiesGroupRequest()
-    pbrequest.kb.uuid = kbid
-    pbrequest.group = group
-    pbrequest.title = item.title or ""
-    pbrequest.color = item.color or ""
-
-    for name, entity in item.add.items():
-        entitypb = pbrequest.add[name]
-        entitypb.value = entity.value
-        entitypb.merged = entity.merged
-        entitypb.represents.extend(entity.represents)
-
-    for name, entity in item.update.items():
-        entitypb = pbrequest.update[name]
-        entitypb.value = entity.value
-        entitypb.merged = entity.merged
-        entitypb.represents.extend(entity.represents)
-
-    pbrequest.delete.extend(item.delete)
-
-    status: UpdateEntitiesGroupResponse = await ingest.UpdateEntitiesGroup(pbrequest)  # type: ignore
-    if status.status == UpdateEntitiesGroupResponse.Status.OK:
-        return
-    elif status.status == UpdateEntitiesGroupResponse.Status.KB_NOT_FOUND:
-        return HTTPNotFound(detail="Knowledge Box does not exist")
-    elif status.status == UpdateEntitiesGroupResponse.Status.ENTITIES_GROUP_NOT_FOUND:
-        return HTTPNotFound(detail="Entities group does not exist")
-    elif status.status == UpdateEntitiesGroupResponse.Status.ERROR:
-        return HTTPInternalServerError(detail="Error on settings entities on a Knowledge box")
-
-
-@api.delete(
-    f"/{KB_PREFIX}/{{kbid}}/entitiesgroup/{{group}}",
-    status_code=200,
-    summary="Delete Knowledge Box Entities",
-    tags=["Knowledge Box Services"],
-    openapi_extra={"x-operation_order": 3},
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def delete_entities(request: Request, kbid: str, group: str):
-    ingest = get_ingest()
-    pbrequest: DelEntitiesRequest = DelEntitiesRequest()
-    pbrequest.kb.uuid = kbid
-    pbrequest.group = group
-
-    status: OpStatusWriter = await ingest.DelEntities(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return None
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(status_code=500, detail="Error on deleting entities from a Knowledge box")
-
-    return Response(status_code=204)
 
 
 @api.post(
@@ -174,7 +44,15 @@ async def delete_entities(request: Request, kbid: str, group: str):
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
-async def set_labelset_endpoint(request: Request, kbid: str, labelset: str, item: LabelSet):
+async def set_labelset_endpoint(
+    request: Request,
+    kbid: str,
+    labelset: str = Path(
+        title="The ID of the labelset to create or update. This is a unique identifier that should be used at search time.",
+        examples=["categories", "movie-genres", "document-types"],
+    ),
+    item: LabelSet = Body(...),
+):
     if item.title is None:
         item.title = labelset