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

nucliadb/openapi.py

@@ -33,11 +33,11 @@ def is_versioned_route(route):
 
 
 def extract_openapi(application, version, commit_id, app_name):
-    app = [
+    app = next(
         route.app
         for route in application.routes
         if is_versioned_route(route) and route.app.version == version
-    ][0]
+    )
     document = get_openapi(
         title=app.title,
         version=app.version,

nucliadb/purge/__init__.py

@@ -19,7 +19,8 @@
 #
 import asyncio
 import importlib.metadata
-from typing import AsyncGenerator
+from collections.abc import AsyncGenerator
+from itertools import batched  # type: ignore
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import NodeError, ShardNotFound
@@ -233,7 +234,7 @@ async def purge_kb_vectorsets(driver: Driver, storage: Storage):
                     fields.extend((await resource.get_fields(force=True)).values())
 
             logger.info(f"Purging {len(fields)} fields for vectorset {vectorset}", extra={"kbid": kbid})
-            for fields_batch in batchify(fields, 20):
+            for fields_batch in batched(fields, n=20):
                 tasks = []
                 for field in fields_batch:
                     if purge_payload.storage_key_kind == VectorSetConfig.StorageKeyKind.UNSET:
@@ -317,9 +318,3 @@ def run() -> int:  # pragma: no cover
     setup_logging()
     errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
     return asyncio.run(main())
-
-
-def batchify(iterable, n=1):
-    """Yield successive n-sized chunks from iterable."""
-    for i in range(0, len(iterable), n):
-        yield iterable[i : i + n]

nucliadb/purge/orphan_shards.py

@@ -20,7 +20,6 @@
 import argparse
 import asyncio
 import importlib.metadata
-from typing import Optional
 
 from grpc.aio import AioRpcError
 from nidx_protos import nodereader_pb2, noderesources_pb2
@@ -113,7 +112,7 @@ async def _get_stored_shards(driver: Driver) -> dict[str, ShardKb]:
     return stored_shards
 
 
-async def _get_kbid(shard_id: str) -> Optional[str]:
+async def _get_kbid(shard_id: str) -> str | None:
     kbid = None
     try:
         req = nodereader_pb2.GetShardRequest()

nucliadb/reader/__init__.py

@@ -19,6 +19,8 @@
 #
 import logging
 
+from fastapi import Header
+
 SERVICE_NAME = "nucliadb.reader"
 logger = logging.getLogger(SERVICE_NAME)
 
@@ -35,3 +37,6 @@ class EndpointFilter(logging.Filter):
 
 # Add filter to the logger
 logging.getLogger("uvicorn.access").addFilter(EndpointFilter())
+
+
+RANGE_HEADER = Header(description="Standard HTTP Range header that enable multipart requests")

nucliadb/reader/api/models.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 TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel
 
@@ -33,14 +33,7 @@ from nucliadb_models.resource import (
 )
 
 if TYPE_CHECKING:  # pragma: no cover
-    ValueType = Optional[
-        Union[
-            models.FieldText,
-            models.FieldFile,
-            models.FieldLink,
-            models.Conversation,
-        ]
-    ]
+    ValueType = models.FieldText | models.FieldFile | models.FieldLink | models.Conversation | None
 else:
     # without Any, pydantic fails to anything as validate() fails using the Union
     ValueType = Any
@@ -50,10 +43,10 @@ class ResourceField(BaseModel):
     field_type: FieldTypeName
     field_id: str
     value: ValueType = None
-    extracted: Optional[ExtractedDataType] = None
-    error: Optional[Error] = None
-    status: Optional[str] = None
-    errors: Optional[list[Error]] = None
+    extracted: ExtractedDataType | None = None
+    error: Error | None = None
+    status: str | None = None
+    errors: list[Error] | None = None
 
 
 FIELD_NAME_TO_EXTRACTED_DATA_FIELD_MAP: dict[FieldTypeName, Any] = {

nucliadb/reader/api/v1/download.py

@@ -18,20 +18,18 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import urllib.parse
-from enum import Enum
-from typing import Optional
+from typing import Annotated
 
 from fastapi import HTTPException
 from fastapi.requests import Request
 from fastapi.responses import Response
 from fastapi_versioning import version
-from starlette.datastructures import Headers
 from starlette.responses import StreamingResponse
 
+from nucliadb.common import datamanagers
 from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR
 from nucliadb.common.models_utils import to_proto
-from nucliadb.ingest.serialize import get_resource_uuid_by_slug
-from nucliadb.reader import SERVICE_NAME, logger
+from nucliadb.reader import RANGE_HEADER, SERVICE_NAME, logger
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_utils.authentication import requires_one
@@ -41,13 +39,8 @@ from nucliadb_utils.utilities import get_storage
 from .router import KB_PREFIX, RESOURCE_PREFIX, RSLUG_PREFIX, api
 
 
-class DownloadType(Enum):
-    EXTRACTED = "extracted"
-    FIELD = "field"
-
-
 @api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/{{field_type}}/{{field_id}}/download/extracted/{{download_field:path}}",  # noqa
+    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/{{field_type}}/{{field_id}}/download/extracted/{{download_field:path}}",
     tags=["Resource fields"],
     status_code=200,
     summary="Download extracted binary file (by slug)",
@@ -61,12 +54,20 @@ async def download_extract_file_rslug_prefix(
     field_type: FieldTypeName,
     field_id: str,
     download_field: str,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
-    return await _download_extract_file(request, kbid, field_type, field_id, download_field, rslug=rslug)
+    return await _download_extract_file(
+        kbid,
+        field_type,
+        field_id,
+        download_field,
+        rslug=rslug,
+        range_request=range,
+    )
 
 
 @api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/{{field_type}}/{{field_id}}/download/extracted/{{download_field:path}}",  # noqa
+    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/{{field_type}}/{{field_id}}/download/extracted/{{download_field:path}}",
     tags=["Resource fields"],
     status_code=200,
     summary="Download extracted binary file (by id)",
@@ -80,18 +81,21 @@ async def download_extract_file_rid_prefix(
     field_type: FieldTypeName,
     field_id: str,
     download_field: str,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
-    return await _download_extract_file(request, kbid, field_type, field_id, download_field, rid=rid)
+    return await _download_extract_file(
+        kbid, field_type, field_id, download_field, rid=rid, range_request=range
+    )
 
 
 async def _download_extract_file(
-    request: Request,
     kbid: str,
     field_type: FieldTypeName,
     field_id: str,
     download_field: str,
-    rid: Optional[str] = None,
-    rslug: Optional[str] = None,
+    rid: str | None = None,
+    rslug: str | None = None,
+    range_request: str | None = None,
 ) -> Response:
     rid = await _get_resource_uuid_from_params(kbid, rid, rslug)
 
@@ -102,7 +106,7 @@ async def _download_extract_file(
 
     sf = storage.file_extracted(kbid, rid, field_type_letter, field_id, download_field)
 
-    return await download_api(sf, request.headers)
+    return await download_api(sf, range_request)
 
 
 @api.get(
@@ -119,8 +123,9 @@ async def download_field_file_rslug_prefix(
     rslug: str,
     field_id: str,
     inline: bool = False,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
-    return await _download_field_file(request, kbid, field_id, rslug=rslug, inline=inline)
+    return await _download_field_file(kbid, field_id, rslug=rslug, range_request=range, inline=inline)
 
 
 @api.get(
@@ -137,16 +142,17 @@ async def download_field_file_rid_prefix(
     rid: str,
     field_id: str,
     inline: bool = False,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
-    return await _download_field_file(request, kbid, field_id, rid=rid, inline=inline)
+    return await _download_field_file(kbid, field_id, rid=rid, range_request=range, inline=inline)
 
 
 async def _download_field_file(
-    request: Request,
     kbid: str,
     field_id: str,
-    rid: Optional[str] = None,
-    rslug: Optional[str] = None,
+    rid: str | None = None,
+    rslug: str | None = None,
+    range_request: str | None = None,
     inline: bool = False,
 ) -> Response:
     rid = await _get_resource_uuid_from_params(kbid, rid, rslug)
@@ -155,11 +161,11 @@ async def _download_field_file(
 
     sf = storage.file_field(kbid, rid, field_id)
 
-    return await download_api(sf, request.headers, inline=inline)
+    return await download_api(sf, range_request=range_request, inline=inline)
 
 
 @api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/conversation/{{field_id}}/download/field/{{message_id}}/{{file_num}}",  # noqa
+    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/conversation/{{field_id}}/download/field/{{message_id}}/{{file_num}}",
     tags=["Resource fields"],
     status_code=200,
     summary="Download conversation binary field (by slug)",
@@ -173,14 +179,20 @@ async def download_field_conversation_rslug_prefix(
     field_id: str,
     message_id: str,
     file_num: int,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
     return await _download_field_conversation_attachment(
-        request, kbid, field_id, message_id, file_num, rslug=rslug
+        kbid,
+        field_id,
+        message_id,
+        file_num,
+        rslug=rslug,
+        range_request=range,
     )
 
 
 @api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/conversation/{{field_id}}/download/field/{{message_id}}/{{file_num}}",  # noqa
+    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/conversation/{{field_id}}/download/field/{{message_id}}/{{file_num}}",
     tags=["Resource fields"],
     status_code=200,
     summary="Download conversation binary field (by id)",
@@ -194,20 +206,26 @@ async def download_field_conversation_attachment_rid_prefix(
     field_id: str,
     message_id: str,
     file_num: int,
+    range: Annotated[str | None, RANGE_HEADER] = None,
 ) -> Response:
     return await _download_field_conversation_attachment(
-        request, kbid, field_id, message_id, file_num, rid=rid
+        kbid,
+        field_id,
+        message_id,
+        file_num,
+        rid=rid,
+        range_request=range,
     )
 
 
 async def _download_field_conversation_attachment(
-    request: Request,
     kbid: str,
     field_id: str,
     message_id: str,
     file_num: int,
-    rid: Optional[str] = None,
-    rslug: Optional[str] = None,
+    rid: str | None = None,
+    rslug: str | None = None,
+    range_request: str | None = None,
 ) -> Response:
     rid = await _get_resource_uuid_from_params(kbid, rid, rslug)
 
@@ -217,11 +235,11 @@ async def _download_field_conversation_attachment(
         kbid, rid, field_id, message_id, attachment_index=file_num
     )
 
-    return await download_api(sf, request.headers)
+    return await download_api(sf, range_request)
 
 
-async def download_api(sf: StorageField, headers: Headers, inline: bool = False):
-    metadata: Optional[ObjectMetadata] = await sf.exists()
+async def download_api(sf: StorageField, range_request: str | None = None, inline: bool = False):
+    metadata: ObjectMetadata | None = await sf.exists()
     if metadata is None:
         raise HTTPException(status_code=404, detail="Specified file doesn't exist")
 
@@ -240,9 +258,8 @@ async def download_api(sf: StorageField, headers: Headers, inline: bool = False)
     }
 
     range = Range()
-    if "range" in headers and file_size > -1:
+    if range_request and file_size > -1:
         status_code = 206
-        range_request = headers["range"]
         try:
             start, end, range_size = parse_media_range(range_request, file_size)
         except NotImplementedError:
@@ -295,13 +312,17 @@ async def download_api(sf: StorageField, headers: Headers, inline: bool = False)
     )
 
 
-async def _get_resource_uuid_from_params(kbid, rid: Optional[str], rslug: Optional[str]) -> str:
+async def _get_resource_uuid_from_params(kbid, rid: str | None, rslug: str | None) -> str:
     if not any([rid, rslug]):
         raise ValueError("Either rid or slug must be set")
 
     if not rid:
         # Attempt to get it from slug
-        rid = await get_resource_uuid_by_slug(kbid, rslug, service_name=SERVICE_NAME)  # type: ignore
+        rid = await datamanagers.atomic.resources.get_resource_uuid_from_slug(
+            kbid=kbid,
+            # mypy doesn't infer that we already checked for slug to be something
+            slug=rslug,  # type: ignore[arg-type]
+        )
         if rid is None:
             raise HTTPException(status_code=404, detail="Resource does not exist")
 

nucliadb/reader/api/v1/export_import.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 AsyncGenerator, AsyncIterable, Union
+from collections.abc import AsyncGenerator, AsyncIterable
 
 from fastapi.responses import StreamingResponse
 from fastapi_versioning import version
@@ -108,7 +108,7 @@ async def download_export_and_delete(
 @version(1)
 async def get_export_status_endpoint(
     request: Request, kbid: str, export_id: str
-) -> Union[StatusResponse, HTTPClientError]:
+) -> StatusResponse | HTTPClientError:
     context = get_app_context(request.app)
     if not await exists_kb(kbid):
         return HTTPClientError(status_code=404, detail="Knowledge Box not found")
@@ -127,7 +127,7 @@ async def get_export_status_endpoint(
 @version(1)
 async def get_import_status_endpoint(
     request: Request, kbid: str, import_id: str
-) -> Union[StatusResponse, HTTPClientError]:
+) -> StatusResponse | HTTPClientError:
     context = get_app_context(request.app)
     if not await exists_kb(kbid):
         return HTTPClientError(status_code=404, detail="Knowledge Box not found")
@@ -137,7 +137,7 @@ async def get_import_status_endpoint(
 
 async def _get_status(
     context: ApplicationContext, type: str, kbid: str, id: str
-) -> Union[StatusResponse, HTTPClientError]:
+) -> StatusResponse | HTTPClientError:
     if type not in ("export", "import"):
         raise ValueError(f"Incorrect type: {type}")
 

nucliadb/reader/api/v1/knowledgebox.py

@@ -17,7 +17,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-from fastapi import HTTPException
+from fastapi import Header, HTTPException
 from fastapi_versioning import version
 from starlette.requests import Request
 
@@ -44,12 +44,20 @@ from nucliadb_utils.authentication import requires, requires_one
 )
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
-async def get_kbs(request: Request, prefix: str = "") -> KnowledgeBoxList:
+async def get_kbs(
+    request: Request,
+    prefix: str = "",
+    x_nucliadb_account: str = Header(default="", include_in_schema=False),
+) -> KnowledgeBoxList:
     driver = get_driver()
     async with driver.ro_transaction() as txn:
         response = KnowledgeBoxList()
         async for kbid, slug in datamanagers.kb.get_kbs(txn, prefix=prefix):
-            response.kbs.append(KnowledgeBoxObjSummary(slug=slug or None, uuid=kbid))
+            response.kbs.append(
+                KnowledgeBoxObjSummary(
+                    slug=user_kb_slug(slug, account_id=x_nucliadb_account) or None, uuid=kbid
+                )
+            )
         return response
 
 
@@ -62,7 +70,9 @@ async def get_kbs(request: Request, prefix: str = "") -> KnowledgeBoxList:
 )
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.READER])
 @version(1)
-async def get_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
+async def get_kb(
+    request: Request, kbid: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
+) -> KnowledgeBoxObj:
     driver = get_driver()
     async with driver.ro_transaction() as txn:
         kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
@@ -71,7 +81,7 @@ async def get_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
 
         return KnowledgeBoxObj(
             uuid=kbid,
-            slug=kb_config.slug,
+            slug=user_kb_slug(kb_config.slug, account_id=x_nucliadb_account),
             config=from_proto.knowledgebox_config(kb_config),
         )
 
@@ -85,12 +95,18 @@ async def get_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
 )
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.READER])
 @version(1)
-async def get_kb_by_slug(request: Request, slug: str) -> KnowledgeBoxObj:
+async def get_kb_by_slug(
+    request: Request, slug: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
+) -> KnowledgeBoxObj:
     driver = get_driver()
     async with driver.ro_transaction() as txn:
-        kbid = await datamanagers.kb.get_kb_uuid(txn, slug=slug)
+        # For cloud, the account id is prepended in order to be able to reuse the same slug in different accounts.
+        kbid = await datamanagers.kb.get_kb_uuid(txn, slug=f"{x_nucliadb_account}:{slug}")
         if kbid is None:
-            raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+            # For onprem, the slug is fully controlled by the user
+            kbid = await datamanagers.kb.get_kb_uuid(txn, slug=slug)
+            if kbid is None:
+                raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
 
         kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
         if kb_config is None:
@@ -98,6 +114,18 @@ async def get_kb_by_slug(request: Request, slug: str) -> KnowledgeBoxObj:
 
         return KnowledgeBoxObj(
             uuid=kbid,
-            slug=kb_config.slug,
+            slug=user_kb_slug(kb_config.slug, account_id=x_nucliadb_account),
             config=from_proto.knowledgebox_config(kb_config),
         )
+
+
+def user_kb_slug(stored_slug: str, account_id: str) -> str:
+    if account_id != "":
+        # On cloud deployments, backend prepends the account id to the user-defined slug.
+        # This is required to make kb slugs reused across different accounts using the same nucliadb.
+        # We strip it so the user does not see it.
+        return stored_slug.split(f"{account_id}:")[-1]
+    else:
+        # On on-prem deployments, the account_id is set to "" by default and we don't need to strip
+        # anything as the backend is not invovled in the kb creation process.
+        return stored_slug

nucliadb/reader/api/v1/learning_config.py

@@ -17,9 +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 Dict
 
-from fastapi import Request
+from fastapi import Header, Request
 from fastapi_versioning import version
 from nuclia_models.config.proto import ExtractConfig, SplitConfiguration
 
@@ -35,7 +34,7 @@ from nucliadb_utils.settings import is_onprem_nucliadb
     path=f"/{KB_PREFIX}/{{kbid}}/models/{{model_id}}/{{filename:path}}",
     status_code=200,
     summary="Download the Knowledege Box model",
-    description="Download the trained model or any other generated file as a result of a training task on a Knowledge Box.",  # noqa
+    description="Download the trained model or any other generated file as a result of a training task on a Knowledge Box.",
     response_model=None,
     tags=["Models"],
 )
@@ -60,15 +59,11 @@ async def download_model(
 )
 @requires_one([NucliaDBRoles.READER, NucliaDBRoles.MANAGER])
 @version(1)
-async def get_configuration(
-    request: Request,
-    kbid: str,
-):
+async def get_configuration(request: Request, kbid: str):
     return await learning_config_proxy(
         request,
         "GET",
         f"/config/{kbid}",
-        extra_headers={"X-STF-USER": request.headers.get("X-NUCLIADB-USER", "")},
     )
 
 
@@ -108,7 +103,6 @@ async def get_model(
         request,
         "GET",
         f"/models/{kbid}/model/{model_id}",
-        extra_headers={"X-STF-USER": request.headers.get("X-NUCLIADB-USER", "")},
     )
 
 
@@ -123,10 +117,35 @@ async def get_model(
 @requires_one([NucliaDBRoles.READER, NucliaDBRoles.MANAGER])
 @version(1)
 async def get_schema_for_configuration_updates(
-    request: Request,
-    kbid: str,
+    request: Request, kbid: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
 ):
-    return await learning_config_proxy(request, "GET", f"/schema/{kbid}")
+    return await learning_config_proxy(
+        request,
+        "GET",
+        f"/schema/{kbid}",
+        headers={"account-id": x_nucliadb_account},
+    )
+
+
+@api.get(
+    path=f"/{KB_PREFIX}/{{kbid}}/generative_providers",
+    status_code=200,
+    summary="Available models for a knowledge box",
+    description="Get all available models for a knowledge box grouped by provider",
+    response_model=None,
+    tags=["Models"],
+)
+@requires_one([NucliaDBRoles.READER, NucliaDBRoles.MANAGER])
+@version(1)
+async def get_models_group_by_providers(
+    request: Request, kbid: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
+):
+    return await learning_config_proxy(
+        request,
+        "GET",
+        f"/generative_providers/{kbid}",
+        headers={"account-id": x_nucliadb_account},
+    )
 
 
 @api.get(
@@ -153,7 +172,7 @@ async def get_schema_for_configuration_creation(
     status_code=200,
     summary="Learning extract strategies",
     description="Get available extract strategies ",
-    response_model=Dict[str, ExtractConfig],
+    response_model=dict[str, ExtractConfig],
     tags=["Extract Strategies"],
 )
 @requires_one([NucliaDBRoles.READER, NucliaDBRoles.MANAGER])
@@ -190,7 +209,7 @@ async def get_extract_strategy_from_id(
     status_code=200,
     summary="Learning split strategies",
     description="Get available split strategies ",
-    response_model=Dict[str, SplitConfiguration],
+    response_model=dict[str, SplitConfiguration],
     tags=["Split Strategies"],
 )
 @requires_one([NucliaDBRoles.READER, NucliaDBRoles.MANAGER])