nucliadb-utils 6.9.6.post5473__py3-none-any.whl → 6.10.0.post5732__py3-none-any.whl

nucliadb_utils/asyncio_utils.py

@@ -19,7 +19,7 @@
 #
 import asyncio
 from collections.abc import Coroutine
-from typing import Any, Optional
+from typing import Any
 
 from nucliadb_utils import logger
 
@@ -30,7 +30,7 @@ class ConcurrentRunner:
     Returns the results of the coroutines in the order they were scheduled.
     """
 
-    def __init__(self, max_tasks: Optional[int] = None):
+    def __init__(self, max_tasks: int | None = None):
         self._tasks: list[asyncio.Task] = []
         self.max_tasks = asyncio.Semaphore(max_tasks) if max_tasks is not None else None
 
@@ -62,7 +62,7 @@ class ConcurrentRunner:
         return results
 
 
-async def run_concurrently(tasks: list[Coroutine], max_concurrent: Optional[int] = None) -> list[Any]:
+async def run_concurrently(tasks: list[Coroutine], max_concurrent: int | None = None) -> list[Any]:
     """
     Runs a list of coroutines concurrently, with a maximum number of tasks running.
     Returns the results of the coroutines in the order they were scheduled.

nucliadb_utils/audit/audit.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 List, Optional
 
 from google.protobuf.timestamp_pb2 import Timestamp
 from nidx_protos.nodereader_pb2 import SearchRequest
@@ -34,14 +33,14 @@ class AuditStorage:
         *,
         kbid: str,
         audit_type: AuditRequest.AuditType.Value,  # type: ignore
-        when: Optional[Timestamp] = None,
-        user: Optional[str] = None,
-        origin: Optional[str] = None,
-        rid: Optional[str] = None,
-        field_metadata: Optional[List[FieldID]] = None,
-        audit_fields: Optional[List[AuditField]] = None,
+        when: Timestamp | None = None,
+        user: str | None = None,
+        origin: str | None = None,
+        rid: str | None = None,
+        field_metadata: list[FieldID] | None = None,
+        audit_fields: list[AuditField] | None = None,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     async def initialize(self):
         pass
@@ -56,10 +55,10 @@ class AuditStorage:
         user: str,
         origin: str,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def send(self, msg: AuditRequest):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def search(
         self,
@@ -70,9 +69,9 @@ class AuditStorage:
         search: SearchRequest,
         timeit: float,
         resources: int,
-        retrieval_rephrased_question: Optional[str] = None,
+        retrieval_rephrased_question: str | None = None,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def chat(
         self,
@@ -81,24 +80,35 @@ class AuditStorage:
         client: int,
         origin: str,
         question: str,
-        rephrased_question: Optional[str],
-        retrieval_rephrased_question: Optional[str],
-        chat_context: List[ChatContext],
-        retrieved_context: List[RetrievedContext],
-        answer: Optional[str],
-        reasoning: Optional[str],
-        learning_id: Optional[str],
+        rephrased_question: str | None,
+        retrieval_rephrased_question: str | None,
+        chat_context: list[ChatContext],
+        retrieved_context: list[RetrievedContext],
+        answer: str | None,
+        reasoning: str | None,
+        learning_id: str | None,
         status_code: int,
-        model: Optional[str],
-        rephrase_time: Optional[float] = None,
-        generative_answer_time: Optional[float] = None,
-        generative_answer_first_chunk_time: Optional[float] = None,
-        generative_reasoning_first_chunk_time: Optional[float] = None,
+        model: str | None,
+        rephrase_time: float | None = None,
+        generative_answer_time: float | None = None,
+        generative_answer_first_chunk_time: float | None = None,
+        generative_reasoning_first_chunk_time: float | None = None,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()
+
+    def retrieve(
+        self,
+        kbid: str,
+        user: str,
+        client: int,
+        origin: str,
+        retrieval_time: float,
+    ):
+        # TODO(decoupled-ask): implement audit for /retrieve
+        ...
 
     def report_storage(self, kbid: str, paragraphs: int, fields: int, bytes: int):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def report_resources(
         self,
@@ -106,10 +116,10 @@ class AuditStorage:
         kbid: str,
         resources: int,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def delete_kb(self, kbid: str):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def feedback(
         self,
@@ -120,7 +130,7 @@ class AuditStorage:
         learning_id: str,
         good: bool,
         task: int,
-        feedback: Optional[str],
-        text_block_id: Optional[str],
+        feedback: str | None,
+        text_block_id: str | None,
     ):
-        raise NotImplementedError
+        raise NotImplementedError()

nucliadb_utils/audit/basic.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 List, Optional
 
 from google.protobuf.timestamp_pb2 import Timestamp
 from nidx_protos.nodereader_pb2 import SearchRequest
@@ -34,19 +33,19 @@ class BasicAuditStorage(AuditStorage):
         self.initialized = True
 
     def message_to_str(self, message: BrokerMessage) -> str:
-        return f"{message.type}+{message.multiid}+{message.audit.user}+{message.kbid}+{message.uuid}+{message.audit.when.ToJsonString()}+{message.audit.origin}+{message.audit.source}"  # noqa
+        return f"{message.type}+{message.multiid}+{message.audit.user}+{message.kbid}+{message.uuid}+{message.audit.when.ToJsonString()}+{message.audit.origin}+{message.audit.source}"
 
     async def report_and_send(
         self,
         *,
         kbid: str,
         audit_type: AuditRequest.AuditType.Value,  # type: ignore
-        when: Optional[Timestamp] = None,
-        user: Optional[str] = None,
-        origin: Optional[str] = None,
-        rid: Optional[str] = None,
-        field_metadata: Optional[List[FieldID]] = None,
-        audit_fields: Optional[List[AuditField]] = None,
+        when: Timestamp | None = None,
+        user: str | None = None,
+        origin: str | None = None,
+        rid: str | None = None,
+        field_metadata: list[FieldID] | None = None,
+        audit_fields: list[AuditField] | None = None,
     ):
         logger.debug(f"AUDIT {audit_type} {kbid} {user} {origin} {rid} {audit_fields}")
 
@@ -71,7 +70,7 @@ class BasicAuditStorage(AuditStorage):
         search: SearchRequest,
         timeit: float,
         resources: int,
-        retrieval_rephrased_question: Optional[str] = None,
+        retrieval_rephrased_question: str | None = None,
     ):
         logger.debug(f"SEARCH {kbid} {user} {origin} ''{search}'' {timeit} {resources}")
 
@@ -82,19 +81,19 @@ class BasicAuditStorage(AuditStorage):
         client_type: int,
         origin: str,
         question: str,
-        rephrased_question: Optional[str],
-        retrieval_rephrased_question: Optional[str],
-        chat_context: List[ChatContext],
-        retrieved_context: List[RetrievedContext],
-        answer: Optional[str],
-        reasoning: Optional[str],
-        learning_id: Optional[str],
+        rephrased_question: str | None,
+        retrieval_rephrased_question: str | None,
+        chat_context: list[ChatContext],
+        retrieved_context: list[RetrievedContext],
+        answer: str | None,
+        reasoning: str | None,
+        learning_id: str | None,
         status_code: int,
-        model: Optional[str],
-        rephrase_time: Optional[float] = None,
-        generative_answer_time: Optional[float] = None,
-        generative_answer_first_chunk_time: Optional[float] = None,
-        generative_reasoning_first_chunk_time: Optional[float] = None,
+        model: str | None,
+        rephrase_time: float | None = None,
+        generative_answer_time: float | None = None,
+        generative_answer_first_chunk_time: float | None = None,
+        generative_reasoning_first_chunk_time: float | None = None,
     ):
         logger.debug(f"CHAT {kbid} {user} {origin}")
 
@@ -121,7 +120,7 @@ class BasicAuditStorage(AuditStorage):
         learning_id: str,
         good: bool,
         task: int,
-        feedback: Optional[str],
-        text_block_id: Optional[str],
+        feedback: str | None,
+        text_block_id: str | None,
     ):
         logger.debug(f"FEEDBACK {kbid} {user} {client_type} {origin}")

nucliadb_utils/audit/stream.py

@@ -21,8 +21,8 @@ import asyncio
 import contextvars
 import json
 import time
+from collections.abc import Callable
 from datetime import datetime, timezone
-from typing import Callable, List, Optional
 
 import backoff
 import mmh3
@@ -71,17 +71,17 @@ class RequestContext:
         self.path: str = ""
 
 
-request_context_var = contextvars.ContextVar[Optional[RequestContext]]("request_context", default=None)
+request_context_var = contextvars.ContextVar[RequestContext | None]("request_context", default=None)
 
 
-def get_trace_id() -> Optional[str]:
+def get_trace_id() -> str | None:
     span = get_current_span()
     if span is None:
         return None
     return format_trace_id(span.get_span_context().trace_id)
 
 
-def get_request_context() -> Optional[RequestContext]:
+def get_request_context() -> RequestContext | None:
     return request_context_var.get()
 
 
@@ -103,7 +103,7 @@ def fill_audit_search_request(audit: AuditSearchRequest, request: SearchRequest)
 
 
 class AuditMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app: ASGIApp, audit_utility_getter: Callable[[], Optional[AuditStorage]]) -> None:
+    def __init__(self, app: ASGIApp, audit_utility_getter: Callable[[], AuditStorage | None]) -> None:
         self.audit_utility_getter = audit_utility_getter
         super().__init__(app)
 
@@ -155,17 +155,17 @@ KB_USAGE_STREAM_SUBJECT = "kb-usage.nuclia_db"
 
 
 class StreamAuditStorage(AuditStorage):
-    task: Optional[asyncio.Task]
+    task: asyncio.Task | None
     initialized: bool
     queue: asyncio.Queue
 
     def __init__(
         self,
-        nats_servers: List[str],
+        nats_servers: list[str],
         nats_target: str,
         partitions: int,
         seed: int,
-        nats_creds: Optional[str] = None,
+        nats_creds: str | None = None,
         service: str = "nucliadb.audit",
     ):
         self.nats_servers = nats_servers
@@ -186,10 +186,10 @@ class StreamAuditStorage(AuditStorage):
 
     async def reconnected_cb(self):
         # See who we are connected to on reconnect.
-        logger.info("Got reconnected to NATS {url}".format(url=self.nc.connected_url))
+        logger.info(f"Got reconnected to NATS {self.nc.connected_url}")
 
     async def error_cb(self, e):
-        logger.error("There was an error connecting to NATS audit: {}".format(e), exc_info=True)
+        logger.error(f"There was an error connecting to NATS audit: {e}", exc_info=True)
 
     async def closed_cb(self):
         logger.info("Connection is closed on NATS")
@@ -269,12 +269,12 @@ class StreamAuditStorage(AuditStorage):
         *,
         kbid: str,
         audit_type: AuditRequest.AuditType.Value,  # type: ignore
-        when: Optional[Timestamp] = None,
-        user: Optional[str] = None,
-        origin: Optional[str] = None,
-        rid: Optional[str] = None,
-        field_metadata: Optional[List[FieldID]] = None,
-        audit_fields: Optional[List[AuditField]] = None,
+        when: Timestamp | None = None,
+        user: str | None = None,
+        origin: str | None = None,
+        rid: str | None = None,
+        field_metadata: list[FieldID] | None = None,
+        audit_fields: list[AuditField] | None = None,
     ):
         auditrequest = AuditRequest()
 
@@ -369,7 +369,7 @@ class StreamAuditStorage(AuditStorage):
         search: SearchRequest,
         timeit: float,
         resources: int,
-        retrieval_rephrased_question: Optional[str] = None,
+        retrieval_rephrased_question: str | None = None,
     ):
         context = get_request_context()
         if context is None:
@@ -420,19 +420,19 @@ class StreamAuditStorage(AuditStorage):
         client_type: int,
         origin: str,
         question: str,
-        rephrased_question: Optional[str],
-        retrieval_rephrased_question: Optional[str],
-        chat_context: List[ChatContext],
-        retrieved_context: List[RetrievedContext],
-        answer: Optional[str],
-        reasoning: Optional[str],
-        learning_id: Optional[str],
+        rephrased_question: str | None,
+        retrieval_rephrased_question: str | None,
+        chat_context: list[ChatContext],
+        retrieved_context: list[RetrievedContext],
+        answer: str | None,
+        reasoning: str | None,
+        learning_id: str | None,
         status_code: int,
-        model: Optional[str],
-        rephrase_time: Optional[float] = None,
-        generative_answer_time: Optional[float] = None,
-        generative_answer_first_chunk_time: Optional[float] = None,
-        generative_reasoning_first_chunk_time: Optional[float] = None,
+        model: str | None,
+        rephrase_time: float | None = None,
+        generative_answer_time: float | None = None,
+        generative_answer_first_chunk_time: float | None = None,
+        generative_reasoning_first_chunk_time: float | None = None,
     ):
         rcontext = get_request_context()
         if rcontext is None:
@@ -482,8 +482,8 @@ class StreamAuditStorage(AuditStorage):
         learning_id: str,
         good: bool,
         task: int,
-        feedback: Optional[str],
-        text_block_id: Optional[str],
+        feedback: str | None,
+        text_block_id: str | None,
     ):
         rcontext = get_request_context()
         if rcontext is None:

nucliadb_utils/authentication.py

@@ -17,12 +17,10 @@
 # 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 functools
 import inspect
 import typing
 from enum import Enum
-from typing import Optional, Tuple
 
 from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
 from starlette.exceptions import HTTPException
@@ -32,7 +30,7 @@ from starlette.websockets import WebSocket
 
 
 class NucliaUser(BaseUser):
-    def __init__(self, username: str, security_groups: Optional[list[str]] = None) -> None:
+    def __init__(self, username: str, security_groups: list[str] | None = None) -> None:
         self.username = username
         self._security_groups = security_groups
 
@@ -45,7 +43,7 @@ class NucliaUser(BaseUser):
         return self.username
 
     @property
-    def security_groups(self) -> Optional[list[str]]:
+    def security_groups(self) -> list[str] | None:
         return self._security_groups
 
 
@@ -63,7 +61,7 @@ class NucliaCloudAuthenticationBackend(AuthenticationBackend):
         self.user_header = user_header
         self.security_groups_header = security_groups_header
 
-    async def authenticate(self, request: HTTPConnection) -> Optional[Tuple[AuthCredentials, BaseUser]]:
+    async def authenticate(self, request: HTTPConnection) -> tuple[AuthCredentials, BaseUser] | None:
         if self.roles_header not in request.headers:
             return None
         else:
@@ -74,9 +72,9 @@ class NucliaCloudAuthenticationBackend(AuthenticationBackend):
         if self.user_header in request.headers:
             user = request.headers[self.user_header]
 
-            raw_security_groups: Optional[str] = request.headers.get(self.security_groups_header)
+            raw_security_groups: str | None = request.headers.get(self.security_groups_header)
 
-            security_groups: Optional[list[str]] = None
+            security_groups: list[str] | None = None
             if raw_security_groups is not None:
                 security_groups = raw_security_groups.split(";")
 
@@ -99,9 +97,9 @@ def has_required_scope(conn: HTTPConnection, scopes: typing.Sequence[str]) -> bo
 
 
 def requires(
-    scopes: typing.Union[str, typing.Sequence[str]],
+    scopes: str | typing.Sequence[str],
     status_code: int = 403,
-    redirect: Optional[str] = None,
+    redirect: str | None = None,
 ) -> typing.Callable:
     # As a fastapi requirement, custom Enum classes have to inherit also from
     # string, so we MUST check for Enum before str
@@ -137,7 +135,7 @@ def requires(
 
             return websocket_wrapper
 
-        elif asyncio.iscoroutinefunction(func):
+        elif inspect.iscoroutinefunction(func):
             # Handle async request/response functions.
             @functools.wraps(func)
             async def async_wrapper(*args: typing.Any, **kwargs: typing.Any) -> Response:

nucliadb_utils/cache/nats.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright (C) 2021 Bosutech XXI S.L.
 #
 # nucliadb is offered under the AGPL v3.0 and as commercial software.
@@ -23,7 +22,6 @@ import functools
 import os
 import uuid
 from inspect import iscoroutinefunction
-from typing import Dict, List, Optional, Union
 
 import nats
 import nats.errors
@@ -54,16 +52,16 @@ async def wait_for_it(future: asyncio.Future, msg):
 class NatsPubsub(PubSubDriver[Msg]):
     _jetstream = None
     _jsm = None
-    _subscriptions: Dict[str, Subscription]
+    _subscriptions: dict[str, Subscription]
     async_callback = True
 
     def __init__(
         self,
         thread: bool = False,
-        name: Optional[str] = "natsutility",
+        name: str | None = "natsutility",
         timeout: float = 2.0,
-        hosts: Optional[List[str]] = None,
-        user_credentials_file: Optional[str] = None,
+        hosts: list[str] | None = None,
+        user_credentials_file: str | None = None,
     ):
         self._hosts = hosts or []
         self._timeout = timeout
@@ -73,11 +71,11 @@ class NatsPubsub(PubSubDriver[Msg]):
         self._uuid = os.environ.get("HOSTNAME", uuid.uuid4().hex)
         self.initialized = False
         self.lock = asyncio.Lock()
-        self.nc: Union[Client, NatsClientTelemetry, None] = None
+        self.nc: Client | NatsClientTelemetry | None = None
         self.user_credentials_file = user_credentials_file
 
     @property
-    def jetstream(self) -> Union[JetStreamContext, JetStreamContextTelemetry]:
+    def jetstream(self) -> JetStreamContext | JetStreamContextTelemetry:
         if self.nc is None:
             raise AttributeError("NC not initialized")
         if self._jetstream is None:
@@ -154,10 +152,10 @@ class NatsPubsub(PubSubDriver[Msg]):
 
     async def reconnected_cb(self):
         # See who we are connected to on reconnect.
-        logger.info("Got reconnected NATS to {url}".format(url=self.nc.connected_url.netloc))
+        logger.info(f"Got reconnected NATS to {self.nc.connected_url.netloc}")
 
     async def error_cb(self, e):
-        logger.info("There was an error connecting to NATS {}".format(e), exc_info=True)
+        logger.info(f"There was an error connecting to NATS {e}", exc_info=True)
 
     async def closed_cb(self):
         logger.info("Connection is closed to NATS")
@@ -173,7 +171,7 @@ class NatsPubsub(PubSubDriver[Msg]):
         else:
             raise ErrConnectionClosed("Could not subscribe")
 
-    async def subscribe(self, handler: Callback, key, group="", subscription_id: Optional[str] = None):
+    async def subscribe(self, handler: Callback, key, group="", subscription_id: str | None = None):
         if subscription_id is None:
             subscription_id = key
 
@@ -189,7 +187,7 @@ class NatsPubsub(PubSubDriver[Msg]):
         else:
             raise ErrConnectionClosed("Could not subscribe")
 
-    async def unsubscribe(self, key: str, subscription_id: Optional[str] = None):
+    async def unsubscribe(self, key: str, subscription_id: str | None = None):
         if subscription_id is None:
             subscription_id = key
 

nucliadb_utils/cache/pubsub.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 Awaitable, Callable, Generic, Optional, TypeVar
+from collections.abc import Awaitable, Callable
+from typing import Generic, TypeVar
 
 T = TypeVar("T")
 
@@ -37,15 +38,15 @@ class PubSubDriver(Generic[T]):
     async def publish(self, channel_name: str, data: bytes):
         raise NotImplementedError()
 
-    async def unsubscribe(self, key: str, subscription_id: Optional[str] = None):
+    async def unsubscribe(self, key: str, subscription_id: str | None = None):
         raise NotImplementedError()
 
     async def subscribe(
         self,
         handler: Callback,
         key: str,
-        group: Optional[str] = None,
-        subscription_id: Optional[str] = None,
+        group: str | None = None,
+        subscription_id: str | None = None,
     ):
         raise NotImplementedError()
 

nucliadb_utils/cache/settings.py

@@ -17,14 +17,13 @@
 # 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 List, Optional
 
 from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):
-    cache_pubsub_nats_url: List[str] = []
-    cache_pubsub_nats_auth: Optional[str] = None
+    cache_pubsub_nats_url: list[str] = []
+    cache_pubsub_nats_auth: str | None = None
 
 
 settings = Settings()

nucliadb_utils/debug.py

@@ -54,7 +54,7 @@ def display_top(snapshot, key_type="lineno", limit=10):  # pragma: no cover
     print(f"Top {limit} lines")
     for index, stat in enumerate(top_stats[:limit], 1):
         frame = stat.traceback[0]
-        print("#%s: %s:%s: %.1f KiB" % (index, frame.filename, frame.lineno, stat.size / 1024))
+        print(f"#{index}: {frame.filename}:{frame.lineno}: {stat.size / 1024:.1f} KiB")
         line = linecache.getline(frame.filename, frame.lineno).strip()
         if line:
             print("    %s" % line)
@@ -62,6 +62,6 @@ def display_top(snapshot, key_type="lineno", limit=10):  # pragma: no cover
     other = top_stats[limit:]
     if other:
         size = sum(stat.size for stat in other)
-        print("%s other: %.1f KiB" % (len(other), size / 1024))
+        print(f"{len(other)} other: {size / 1024:.1f} KiB")
     total = sum(stat.size for stat in top_stats)
     print("Total allocated size: %.1f KiB" % (total / 1024))

nucliadb_utils/encryption/settings.py

@@ -16,14 +16,13 @@
 #
 # 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 pydantic import Field
 from pydantic_settings import BaseSettings
 
 
 class EncryptionSettings(BaseSettings):
-    encryption_secret_key: Optional[str] = Field(
+    encryption_secret_key: str | None = Field(
         default=None,
         title="Encryption Secret Key",
         description="""Secret key used for encryption and decryption of sensitive data in the database.

nucliadb_utils/fastapi/openapi.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 Union
 
 from fastapi.applications import FastAPI
 from starlette.routing import Mount
@@ -34,7 +33,7 @@ def format_scopes(scope_list):
     return "\n".join(f"- `{scope}`" for scope in scope_list)
 
 
-def extend_openapi(app: Union[FastAPI, Mount]):  # pragma: no cover
+def extend_openapi(app: FastAPI | Mount):  # pragma: no cover
     for route in app.routes:
         # mypy complains about BaseRoute not having endpoint and
         # description attributes, but routes passed here always have