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

nucliadb/standalone/auth.py

@@ -19,7 +19,6 @@
 import base64
 import logging
 import time
-from typing import Optional
 
 import orjson
 from jwcrypto import jwe, jwk  # type: ignore
@@ -51,7 +50,7 @@ def get_mapped_roles(*, settings: Settings, data: dict[str, str]) -> list[str]:
 
 async def authenticate_auth_token(
     settings: Settings, request: HTTPConnection
-) -> Optional[tuple[AuthCredentials, BaseUser]]:
+) -> tuple[AuthCredentials, BaseUser] | None:
     if "eph-token" not in request.query_params or settings.jwk_key is None:
         return None
 
@@ -81,7 +80,7 @@ class AuthHeaderAuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -109,7 +108,7 @@ class OAuth2AuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -160,7 +159,7 @@ class BasicAuthAuthenticationBackend(NucliaCloudAuthenticationBackend):
     def __init__(self, settings: Settings) -> None:
         self.settings = settings
 
-    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp
@@ -189,7 +188,7 @@ class UpstreamNaiveAuthenticationBackend(NucliaCloudAuthenticationBackend):
             user_header=settings.auth_policy_user_header,
         )
 
-    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
         token_resp = await authenticate_auth_token(self.settings, request)
         if token_resp is not None:
             return token_resp

nucliadb/standalone/lifecycle.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/>.
 #
-import asyncio
+import inspect
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
@@ -56,7 +56,7 @@ async def lifespan(app: FastAPI):
         yield
 
         for finalizer in SYNC_FINALIZERS:
-            if asyncio.iscoroutinefunction(finalizer):
+            if inspect.iscoroutinefunction(finalizer):
                 await finalizer()
             else:
                 finalizer()

nucliadb/standalone/run.py

@@ -21,7 +21,6 @@ import asyncio
 import logging
 import os
 import sys
-from typing import Optional
 
 import argdantic
 import uvicorn  # type: ignore
@@ -148,9 +147,8 @@ def run():
     server.run()
 
 
-def get_latest_nucliadb() -> Optional[str]:
-    loop = asyncio.get_event_loop()
-    return loop.run_until_complete(versions.latest_nucliadb())
+def get_latest_nucliadb() -> str | None:
+    return asyncio.run(versions.latest_nucliadb())
 
 
 async def run_async_nucliadb(settings: Settings) -> uvicorn.Server:

nucliadb/standalone/settings.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from enum import Enum
-from typing import Optional
 
 import pydantic
 
@@ -44,11 +43,11 @@ class Settings(DriverSettings, StorageSettings, ExtendedStorageSettings):
     # all settings here are mapped in to other env var settings used
     # in the app. These are helper settings to make things easier to
     # use with standalone app vs cluster app.
-    nua_api_key: Optional[str] = pydantic.Field(
+    nua_api_key: str | None = pydantic.Field(
         default=None,
-        description="Nuclia Understanding API Key. Read how to generate a NUA Key here: https://docs.nuclia.dev/docs/rag/advanced/understanding/intro#get-a-nua-key",  # noqa
+        description="Nuclia Understanding API Key. Read how to generate a NUA Key here: https://docs.nuclia.dev/docs/rag/advanced/understanding/intro#get-a-nua-key",
     )
-    zone: Optional[str] = pydantic.Field(default=None, description="Nuclia Understanding API Zone ID")
+    zone: str | None = pydantic.Field(default=None, description="Nuclia Understanding API Zone ID")
     http_host: str = pydantic.Field(default="0.0.0.0", description="HTTP Port")
     http_port: int = pydantic.Field(default=8080, description="HTTP Port")
     ingest_grpc_port: int = pydantic.Field(default=8030, description="Ingest GRPC Port")
@@ -83,7 +82,7 @@ class Settings(DriverSettings, StorageSettings, ExtendedStorageSettings):
         description="Default role to assign to user that is authenticated \
                     upstream. Not used with `upstream_naive` auth policy.",
     )
-    auth_policy_role_mapping: Optional[dict[str, dict[str, list[NucliaDBRoles]]]] = pydantic.Field(
+    auth_policy_role_mapping: dict[str, dict[str, list[NucliaDBRoles]]] | None = pydantic.Field(
         default=None,
         description="""
 Role mapping for `upstream_auth_header`, `upstream_oauth2` and `upstream_basicauth` auth policies.
@@ -97,7 +96,7 @@ Examples:
 """,
     )
 
-    jwk_key: Optional[str] = pydantic.Field(
+    jwk_key: str | None = pydantic.Field(
         default=None,
         description="JWK key used for temporary token generation and validation.",
     )

nucliadb/standalone/versions.py

@@ -20,7 +20,6 @@
 import enum
 import importlib.metadata
 import logging
-from typing import Optional
 
 from cachetools import TTLCache
 
@@ -45,11 +44,11 @@ def installed_nucliadb() -> str:
     return get_installed_version(StandalonePackages.NUCLIADB.value)
 
 
-async def latest_nucliadb() -> Optional[str]:
+async def latest_nucliadb() -> str | None:
     return await get_latest_version(StandalonePackages.NUCLIADB.value)
 
 
-def nucliadb_updates_available(installed: str, latest: Optional[str]) -> bool:
+def nucliadb_updates_available(installed: str, latest: str | None) -> bool:
     if latest is None:
         return False
     return is_newer_release(installed, latest)
@@ -96,7 +95,7 @@ def get_installed_version(package_name: str) -> str:
     return importlib.metadata.distribution(package_name).version
 
 
-async def get_latest_version(package: str) -> Optional[str]:
+async def get_latest_version(package: str) -> str | None:
     result = CACHE.get(package, None)
     if result is None:
         try:

nucliadb/tasks/consumer.py

@@ -19,9 +19,10 @@
 #
 
 import asyncio
-from typing import Generic, Optional, Type
+from typing import Generic
 
 import nats
+import nats.js.api
 import pydantic
 from nats.aio.client import Msg
 
@@ -43,8 +44,9 @@ class NatsTaskConsumer(Generic[MsgType]):
         stream: NatsStream,
         consumer: NatsConsumer,
         callback: Callback,
-        msg_type: Type[MsgType],
-        max_concurrent_messages: Optional[int] = None,
+        msg_type: type[MsgType],
+        max_concurrent_messages: int | None = None,
+        max_deliver: int | None = None,
     ):
         self.name = name
         self.stream = stream
@@ -52,6 +54,7 @@ class NatsTaskConsumer(Generic[MsgType]):
         self.callback = callback
         self.msg_type = msg_type
         self.max_concurrent_messages = max_concurrent_messages
+        self.max_deliver = max_deliver
         self.initialized = False
         self.running_tasks: list[asyncio.Task] = []
         self.subscription = None
@@ -71,7 +74,8 @@ class NatsTaskConsumer(Generic[MsgType]):
         for task in self.running_tasks:
             task.cancel()
         try:
-            await asyncio.wait(self.running_tasks, timeout=5)
+            if len(self.running_tasks) > 0:
+                await asyncio.wait(self.running_tasks, timeout=5)
             self.running_tasks.clear()
         except asyncio.TimeoutError:
             pass
@@ -96,6 +100,7 @@ class NatsTaskConsumer(Generic[MsgType]):
                 ack_wait=nats_consumer_settings.nats_ack_wait,
                 idle_heartbeat=nats_consumer_settings.nats_idle_heartbeat,
                 max_ack_pending=max_ack_pending,
+                max_deliver=self.max_deliver,
             ),
         )
         logger.info(
@@ -168,8 +173,6 @@ class NatsTaskConsumer(Generic[MsgType]):
                     },
                 )
                 await msg.ack()
-            finally:
-                return
 
 
 def create_consumer(
@@ -177,8 +180,9 @@ def create_consumer(
     stream: NatsStream,
     consumer: NatsConsumer,
     callback: Callback,
-    msg_type: Type[MsgType],
-    max_concurrent_messages: Optional[int] = None,
+    msg_type: type[MsgType],
+    max_concurrent_messages: int | None = None,
+    max_retries: int = 100,
 ) -> NatsTaskConsumer[MsgType]:
     """
     Returns a non-initialized consumer
@@ -190,4 +194,5 @@ def create_consumer(
         callback=callback,
         msg_type=msg_type,
         max_concurrent_messages=max_concurrent_messages,
+        max_deliver=max_retries,
     )

nucliadb/tasks/models.py

@@ -17,7 +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 typing import Any, Callable, Coroutine, TypeVar
+from collections.abc import Callable, Coroutine
+from typing import Any, TypeVar
 
 import pydantic
 

nucliadb/tasks/producer.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 Generic, Type
+from typing import Generic
 
 from nucliadb.tasks.logger import logger
 from nucliadb.tasks.models import MsgType
@@ -32,7 +32,7 @@ class NatsTaskProducer(Generic[MsgType]):
         name: str,
         stream: NatsStream,
         producer_subject: str,
-        msg_type: Type[MsgType],
+        msg_type: type[MsgType],
     ):
         self.name = name
         self.stream = stream
@@ -69,7 +69,7 @@ def create_producer(
     name: str,
     stream: NatsStream,
     producer_subject: str,
-    msg_type: Type[MsgType],
+    msg_type: type[MsgType],
 ) -> NatsTaskProducer[MsgType]:
     """
     Returns a non-initialized producer.

nucliadb/tasks/retries.py

@@ -19,9 +19,10 @@
 #
 import functools
 import logging
+from collections.abc import Callable
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Callable, Optional, cast
+from typing import cast
 
 from pydantic import BaseModel
 
@@ -44,7 +45,7 @@ class TaskMetadata(BaseModel):
     status: Status
     retries: int = 0
     error_messages: list[str] = []
-    last_modified: Optional[datetime] = None
+    last_modified: datetime | None = None
 
 
 class TaskRetryHandler:
@@ -87,7 +88,7 @@ class TaskRetryHandler:
             kbid=self.kbid, task_type=self.task_type, task_id=self.task_id
         )
 
-    async def get_metadata(self) -> Optional[TaskMetadata]:
+    async def get_metadata(self) -> TaskMetadata | None:
         return await _get_metadata(self.context.kv_driver, self.metadata_key)
 
     async def set_metadata(self, metadata: TaskMetadata) -> None:
@@ -150,7 +151,7 @@ class TaskRetryHandler:
         return wrapper
 
 
-async def _get_metadata(kv_driver: Driver, metadata_key: str) -> Optional[TaskMetadata]:
+async def _get_metadata(kv_driver: Driver, metadata_key: str) -> TaskMetadata | None:
     async with kv_driver.ro_transaction() as txn:
         metadata = await txn.get(metadata_key)
         if metadata is None:
@@ -173,7 +174,7 @@ async def purge_metadata(kv_driver: Driver) -> int:
         return 0
 
     total_purged = 0
-    start: Optional[str] = ""
+    start: str | None = ""
     while True:
         start, purged = await purge_batch(kv_driver, start)
         total_purged += purged
@@ -183,8 +184,8 @@ async def purge_metadata(kv_driver: Driver) -> int:
 
 
 async def purge_batch(
-    kv_driver: PGDriver, start: Optional[str] = None, batch_size: int = 200
-) -> tuple[Optional[str], int]:
+    kv_driver: PGDriver, start: str | None = None, batch_size: int = 200
+) -> tuple[str | None, int]:
     """
     Returns the next start key and the number of purged records. If start is None, it means there are no more records to purge.
     """

nucliadb/train/api/utils.py

@@ -19,12 +19,10 @@
 #
 
 
-from typing import Optional
-
 from nucliadb.train.utils import get_shard_manager
 
 
-async def get_kb_partitions(kbid: str, prefix: Optional[str] = None) -> list[str]:
+async def get_kb_partitions(kbid: str, prefix: str | None = None) -> list[str]:
     shard_manager = get_shard_manager()
     shards = await shard_manager.get_shards_by_kbid_inner(kbid=kbid)
     valid_shards = []

nucliadb/train/api/v1/shards.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import json
-from typing import Optional
 
 import google.protobuf.message
 import pydantic
@@ -63,7 +62,7 @@ async def object_get_response(
     )
 
 
-async def get_trainset(request: Request) -> tuple[TrainSet, Optional[FilterExpression]]:
+async def get_trainset(request: Request) -> tuple[TrainSet, FilterExpression | None]:
     if request.headers.get("Content-Type") == "application/json":
         try:
             trainset_model = TrainSetModel.model_validate(await request.json())

nucliadb/train/api/v1/trainset.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
 
 from fastapi import HTTPException, Request
 from fastapi_versioning import version
@@ -57,7 +56,7 @@ async def get_partitions_prefix(request: Request, kbid: str, prefix: str) -> Tra
     return await get_partitions(kbid, prefix=prefix)
 
 
-async def get_partitions(kbid: str, prefix: Optional[str] = None) -> TrainSetPartitions:
+async def get_partitions(kbid: str, prefix: str | None = None) -> TrainSetPartitions:
     try:
         all_keys = await get_kb_partitions(kbid, prefix)
     except ShardNotFound:

nucliadb/train/app.py

@@ -50,7 +50,6 @@ errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
 fastapi_settings = dict(
     debug=running_settings.debug,
-    middleware=middleware,
     lifespan=lifespan,
     exception_handlers={
         Exception: global_exception_handler,
@@ -71,6 +70,7 @@ application = VersionedFastAPI(
     prefix_format=f"/{API_PREFIX}/v{{major}}",
     default_version=(1, 0),
     enable_latest=False,
+    middleware=middleware,
     kwargs=fastapi_settings,
 )
 

nucliadb/train/generator.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, Callable, Optional
+from collections.abc import AsyncIterator, Callable
 
 from fastapi import HTTPException
 from grpc import StatusCode
@@ -53,11 +53,11 @@ from nucliadb.train.utils import get_shard_manager
 from nucliadb_models.filters import FilterExpression
 from nucliadb_protos.dataset_pb2 import TaskType, TrainSet
 
-BatchGenerator = Callable[[str, TrainSet, str, Optional[FilterExpression]], AsyncIterator[TrainBatch]]
+BatchGenerator = Callable[[str, TrainSet, str, FilterExpression | None], AsyncIterator[TrainBatch]]
 
 
 async def generate_train_data(
-    kbid: str, shard: str, trainset: TrainSet, filter_expression: Optional[FilterExpression] = None
+    kbid: str, shard: str, trainset: TrainSet, filter_expression: FilterExpression | None = None
 ):
     # Get the data structure to generate data
     shard_manager = get_shard_manager()
@@ -66,7 +66,7 @@ async def generate_train_data(
     if trainset.batch_size == 0:
         trainset.batch_size = 50
 
-    batch_generator: Optional[BatchGenerator] = None
+    batch_generator: BatchGenerator | None = None
 
     if trainset.type == TaskType.FIELD_CLASSIFICATION:
         batch_generator = field_classification_batch_generator

nucliadb/train/generators/field_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 nidx_protos.nodereader_pb2 import StreamRequest
 
@@ -39,7 +39,7 @@ def field_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[FieldClassificationBatch, None]:
     generator = generate_field_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, FieldClassificationBatch)

nucliadb/train/generators/field_streaming.py

@@ -19,7 +19,7 @@
 #
 
 import asyncio
-from typing import AsyncGenerator, AsyncIterable, Optional
+from collections.abc import AsyncGenerator, AsyncIterable
 
 from nidx_protos.nodereader_pb2 import DocumentItem, StreamRequest
 
@@ -45,7 +45,7 @@ def field_streaming_batch_generator(
     kbid: str,
     trainset: TrainSet,
     shard_replica_id: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
 ) -> AsyncGenerator[FieldStreamingBatch, None]:
     generator = generate_field_streaming_payloads(kbid, trainset, shard_replica_id, filter_expression)
     batch_generator = batchify(generator, trainset.batch_size, FieldStreamingBatch)
@@ -53,7 +53,7 @@ def field_streaming_batch_generator(
 
 
 async def generate_field_streaming_payloads(
-    kbid: str, trainset: TrainSet, shard_replica_id: str, filter_expression: Optional[FilterExpression]
+    kbid: str, trainset: TrainSet, shard_replica_id: str, filter_expression: FilterExpression | None
 ) -> AsyncGenerator[FieldSplitData, None]:
     request = StreamRequest()
     request.shard_id.id = shard_replica_id
@@ -192,7 +192,7 @@ async def _fetch_basic(kbid: str, fsd: FieldSplitData):
         fsd.basic.CopyFrom(basic)
 
 
-async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> Optional[ExtractedText]:
+async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> ExtractedText | None:
     orm_resource = await get_resource_from_cache_or_db(kbid, rid)
 
     if orm_resource is None:
@@ -208,7 +208,7 @@ async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> Op
 
 async def get_field_metadata(
     kbid: str, rid: str, field: str, field_type: str
-) -> Optional[FieldComputedMetadata]:
+) -> FieldComputedMetadata | None:
     orm_resource = await get_resource_from_cache_or_db(kbid, rid)
 
     if orm_resource is None:
@@ -222,7 +222,7 @@ async def get_field_metadata(
     return field_metadata
 
 
-async def get_field_basic(kbid: str, rid: str, field: str, field_type: str) -> Optional[Basic]:
+async def get_field_basic(kbid: str, rid: str, field: str, field_type: str) -> Basic | None:
     orm_resource = await get_resource_from_cache_or_db(kbid, rid)
 
     if orm_resource is None:

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(

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 = []