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

nucliadb/search/search/summarize.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from typing import Optional
 
 from nucliadb.common import datamanagers
 from nucliadb.common.maindb.utils import get_driver
@@ -36,7 +35,7 @@ from nucliadb_models.search import (
 from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_utils.utilities import get_storage
 
-ExtractedTexts = list[tuple[str, str, Optional[ExtractedText]]]
+ExtractedTexts = list[tuple[str, str, ExtractedText | None]]
 
 MAX_GET_EXTRACTED_TEXT_OPS = 20
 
@@ -46,7 +45,7 @@ class NoResourcesToSummarize(Exception):
 
 
 async def summarize(
-    kbid: str, request: SummarizeRequest, extra_predict_headers: Optional[dict[str, str]]
+    kbid: str, request: SummarizeRequest, extra_predict_headers: dict[str, str] | None
 ) -> SummarizedResponse:
     predict_request = SummarizeModel()
     predict_request.generative_model = request.generative_model
@@ -87,7 +86,7 @@ async def get_extracted_texts(kbid: str, resource_uuids_or_slugs: list[str]) ->
             if uuid is None:
                 logger.warning(f"Resource {uuid_or_slug} not found in KB", extra={"kbid": kbid})
                 continue
-            resource_orm = Resource(txn=txn, storage=storage, kb=kb_orm, uuid=uuid)
+            resource_orm = Resource(txn=txn, storage=storage, kbid=kbid, uuid=uuid)
             fields = await resource_orm.get_fields(force=True)
             for _, field in fields.items():
                 task = asyncio.create_task(get_extracted_text(uuid_or_slug, field, max_tasks))
@@ -115,14 +114,14 @@ async def get_extracted_texts(kbid: str, resource_uuids_or_slugs: list[str]) ->
 
 async def get_extracted_text(
     uuid_or_slug, field: Field, max_operations: asyncio.Semaphore
-) -> tuple[str, str, Optional[ExtractedText]]:
+) -> tuple[str, str, ExtractedText | None]:
     async with max_operations:
         extracted_text = await field.get_extracted_text(force=True)
         field_key = f"{field.type}/{field.id}"
         return uuid_or_slug, field_key, extracted_text
 
 
-async def get_resource_uuid(kbobj: KnowledgeBox, uuid_or_slug: str) -> Optional[str]:
+async def get_resource_uuid(kbobj: KnowledgeBox, uuid_or_slug: str) -> str | None:
     """
     Return the uuid of the resource with the given uuid_or_slug.
     """

nucliadb/search/search/utils.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import Optional
 
 from pydantic import BaseModel
 
@@ -30,7 +29,7 @@ from nucliadb_utils.utilities import has_feature
 logger = logging.getLogger(__name__)
 
 
-async def filter_hidden_resources(kbid: str, show_hidden: bool) -> Optional[bool]:
+async def filter_hidden_resources(kbid: str, show_hidden: bool) -> bool | None:
     kb_config = await kb.get_config(kbid=kbid)
     hidden_enabled = kb_config and kb_config.hidden_resources_enabled
     if hidden_enabled and not show_hidden:
@@ -41,8 +40,8 @@ async def filter_hidden_resources(kbid: str, show_hidden: bool) -> Optional[bool
 
 def min_score_from_query_params(
     min_score_bm25: float,
-    min_score_semantic: Optional[float],
-    deprecated_min_score: Optional[float],
+    min_score_semantic: float | None,
+    deprecated_min_score: float | None,
 ) -> MinScore:
     # Keep backward compatibility with the deprecated min_score parameter
     semantic = deprecated_min_score if min_score_semantic is None else min_score_semantic

nucliadb/search/settings.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
 
 from pydantic import Field
 
@@ -43,7 +42,7 @@ class Settings(DriverSettings):
         title="Prequeries max parallel",
         description="The maximum number of prequeries to run in parallel per /ask request",
     )
-    nidx_address: Optional[str] = Field(default=None)
+    nidx_address: str | None = Field(default=None)
 
 
 settings = Settings()

nucliadb/standalone/api_router.py

@@ -57,7 +57,7 @@ async def api_config_check(request: Request):
                     valid_nua_key = True
             except Exception as exc:
                 logger.warning(f"Error validating nua key", exc_info=exc)
-                nua_key_check_error = f"Error checking NUA key: {str(exc)}"
+                nua_key_check_error = f"Error checking NUA key: {exc!s}"
     return JSONResponse(
         {
             "nua_api_key": {

nucliadb/standalone/app.py

@@ -31,7 +31,7 @@ from starlette.responses import HTMLResponse
 from starlette.routing import Mount
 
 import nucliadb_admin_assets  # type: ignore
-from nucliadb.middleware import ProcessTimeHeaderMiddleware
+from nucliadb.middleware import ClientErrorPayloadLoggerMiddleware, ProcessTimeHeaderMiddleware
 from nucliadb.reader import API_PREFIX
 from nucliadb.reader.api.v1.router import api as api_reader_v1
 from nucliadb.search.api.v1.router import api as api_search_v1
@@ -79,7 +79,7 @@ HOMEPAGE_HTML = """
     </ul>
 </body>
 </html>
-"""  # noqa: E501
+"""
 
 
 def application_factory(settings: Settings) -> FastAPI:
@@ -95,13 +95,13 @@ def application_factory(settings: Settings) -> FastAPI:
             backend=get_auth_backend(settings),
         ),
         Middleware(AuditMiddleware, audit_utility_getter=get_audit),
+        Middleware(ClientErrorPayloadLoggerMiddleware),
     ]
     if running_settings.debug:
         middleware.append(Middleware(ProcessTimeHeaderMiddleware))
 
     fastapi_settings = dict(
         debug=running_settings.debug,
-        middleware=middleware,
         lifespan=lifespan,
         exception_handlers={
             Exception: global_exception_handler,
@@ -122,6 +122,7 @@ def application_factory(settings: Settings) -> FastAPI:
         prefix_format=f"/{API_PREFIX}/v{{major}}",
         default_version=(1, 0),
         enable_latest=False,
+        middleware=middleware,
         kwargs=fastapi_settings,
     )
 

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
@@ -116,6 +115,9 @@ def run():
     if nuclia_settings.nuclia_service_account:
         settings_to_output["NUA API key"] = "Configured ✔"
         settings_to_output["NUA API zone"] = nuclia_settings.nuclia_zone
+        settings_to_output["NUA API url"] = (
+            nuclia_settings.nuclia_public_url.format(zone=nuclia_settings.nuclia_zone) + "/api"
+        )
 
     settings_to_output_fmted = "\n".join(
         [f"||      - {k}:{' ' * (27 - len(k))}{v}" for k, v in settings_to_output.items()]
@@ -145,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: