nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/writer/tus/s3.py

@@ -19,21 +19,22 @@
 #
 from __future__ import annotations
 
+import base64
 import uuid
 from contextlib import AsyncExitStack
-from typing import AsyncIterator, Optional
+from typing import Optional
 
 import aiobotocore  # type: ignore
 import aiohttp
-import backoff  # type: ignore
+import backoff
 import botocore  # type: ignore
 from aiobotocore.session import AioSession  # type: ignore
-from nucliadb_protos.resources_pb2 import CloudFile
 
 from nucliadb.writer import logger
 from nucliadb.writer.tus.dm import FileDataManager
-from nucliadb.writer.tus.exceptions import CloudFileNotFound, ResumableURINotAvailable
+from nucliadb.writer.tus.exceptions import ResumableURINotAvailable
 from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
+from nucliadb_protos.resources_pb2 import CloudFile
 from nucliadb_utils.storages.s3 import (
     CHUNK_SIZE,
     MIN_UPLOAD_SIZE,
@@ -53,9 +54,7 @@ class S3FileStorageManager(FileStorageManager):
     chunk_size = CHUNK_SIZE
     min_upload_size = MIN_UPLOAD_SIZE
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def _abort_multipart(self, dm: FileDataManager):
         try:
             mpu = dm.get("mpu")
@@ -72,21 +71,25 @@ class S3FileStorageManager(FileStorageManager):
         if dm.get("mpu") is not None:
             await self._abort_multipart(dm)
 
+        custom_metadata: dict[str, str] = {
+            "base64_filename": base64.b64encode((dm.filename or "").encode()).decode(),
+            "content_type": dm.content_type or "",
+            "size": str(dm.size),
+        }
+
         await dm.update(
             path=path,
             upload_file_id=upload_file_id,
             multipart={"Parts": []},
             block=1,
-            mpu=await self._create_multipart(path, bucket),
+            mpu=await self._create_multipart(path, bucket, custom_metadata),
             bucket=bucket,
         )
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
-    async def _create_multipart(self, path, bucket):
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
+    async def _create_multipart(self, path, bucket, custom_metadata: dict[str, str]):
         return await self.storage._s3aioclient.create_multipart_upload(
-            Bucket=bucket, Key=path
+            Bucket=bucket, Key=path, Metadata=custom_metadata
         )
 
     async def append(self, dm: FileDataManager, iterable, offset) -> int:
@@ -96,16 +99,12 @@ class S3FileStorageManager(FileStorageManager):
             size += len(chunk)
             part = await self._upload_part(dm, chunk)
             multipart = dm.get("multipart")
-            multipart["Parts"].append(
-                {"PartNumber": dm.get("block"), "ETag": part["ETag"]}
-            )
+            multipart["Parts"].append({"PartNumber": dm.get("block"), "ETag": part["ETag"]})
             await dm.update(multipart=multipart, block=dm.get("block") + 1)
 
         return size
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def _upload_part(self, dm: FileDataManager, data):
         mpu = dm.get("mpu")
         if mpu is None:
@@ -128,18 +127,14 @@ class S3FileStorageManager(FileStorageManager):
         await dm.finish()
         return path
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def _complete_multipart_upload(self, dm: FileDataManager):
         # if blocks is 0, it means the file is of zero length so we need to
         # trick it to finish a multiple part with no data.
         if dm.get("block") == 1:
             part = await self._upload_part(dm, b"")
             multipart = dm.get("multipart")
-            multipart["Parts"].append(
-                {"PartNumber": dm.get("block"), "ETag": part["ETag"]}
-            )
+            multipart["Parts"].append({"PartNumber": dm.get("block"), "ETag": part["ETag"]})
             await dm.update(multipart=multipart, block=dm.get("block") + 1)
         await self.storage._s3aioclient.complete_multipart_upload(
             Bucket=dm.get("bucket"),
@@ -148,45 +143,10 @@ class S3FileStorageManager(FileStorageManager):
             MultipartUpload=dm.get("multipart"),
         )
 
-    @backoff.on_exception(
-        backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3
-    )
+    @backoff.on_exception(backoff.expo, RETRIABLE_EXCEPTIONS, jitter=backoff.random_jitter, max_tries=3)
     async def _download(self, uri: str, kbid: str, **kwargs):
         bucket = self.storage.get_bucket_name(kbid)
-        return await self.storage._s3aioclient.get_object(
-            Bucket=bucket, Key=uri, **kwargs
-        )
-
-    async def iter_data(
-        self, uri: str, kbid: str, headers: Optional[dict[str, str]] = None
-    ):
-        if headers is None:
-            headers = {}
-        try:
-            downloader = await self._download(uri, kbid, **headers)
-        except self.storage._s3aioclient.exceptions.NoSuchKey:
-            raise CloudFileNotFound()
-
-        # we do not want to timeout ever from this...
-        # downloader['Body'].set_socket_timeout(999999)
-        stream = downloader["Body"]
-        data = await stream.read(CHUNK_SIZE)
-        while True:
-            if not data:
-                break
-            yield data
-            data = await stream.read(CHUNK_SIZE)
-
-    async def read_range(
-        self, uri, kbid: str, start: int, end: int
-    ) -> AsyncIterator[bytes]:
-        """
-        Iterate through ranges of data
-        """
-        async for chunk in self.iter_data(
-            uri, kbid, headers={"Range": f"bytes={start}-{end - 1}"}
-        ):
-            yield chunk
+        return await self.storage._s3aioclient.get_object(Bucket=bucket, Key=uri, **kwargs)
 
     async def delete_upload(self, uri: str, kbid: str):
         bucket = self.storage.get_bucket_name(kbid)
@@ -198,6 +158,10 @@ class S3FileStorageManager(FileStorageManager):
         else:
             raise AttributeError("No valid uri")
 
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        if uploaded_bytes % self.min_upload_size != 0:
+            raise ValueError(f"Intermediate chunks need to be multiples of {self.min_upload_size} bytes")
+
 
 class S3BlobStore(BlobStore):
     async def check_exists(self, bucket_name: str) -> bool:
@@ -213,9 +177,7 @@ class S3BlobStore(BlobStore):
     async def create_bucket(self, bucket):
         exists = await self.check_exists(bucket)
         if not exists:
-            await create_bucket(
-                self._s3aioclient, bucket, self.bucket_tags, self.region_name
-            )
+            await create_bucket(self._s3aioclient, bucket, self.bucket_tags, self.region_name)
         return exists
 
     async def finalize(self):
@@ -247,9 +209,7 @@ class S3BlobStore(BlobStore):
             verify=verify_ssl,
             use_ssl=ssl,
             region_name=region_name,
-            config=aiobotocore.config.AioConfig(
-                None, max_pool_connections=max_pool_connections
-            ),
+            config=aiobotocore.config.AioConfig(None, max_pool_connections=max_pool_connections),
         )
         session = AioSession()
         self._s3aioclient = await self._exit_stack.enter_async_context(

nucliadb/writer/tus/storage.py

@@ -21,15 +21,8 @@ from __future__ import annotations
 
 from typing import AsyncIterator, Optional
 
-from lru import LRU  # type: ignore
-from nucliadb_protos.resources_pb2 import CloudFile
-from starlette.responses import StreamingResponse
-
-from nucliadb.writer import logger
 from nucliadb.writer.tus.dm import FileDataManager
-from nucliadb.writer.tus.exceptions import HTTPRangeNotSatisfiable
-
-CACHED_BUCKETS = LRU(50)  # type: ignore
+from nucliadb_protos.resources_pb2 import CloudFile
 
 
 class BlobStore:
@@ -56,14 +49,9 @@ class FileStorageManager:
     chunk_size: int
     min_upload_size: Optional[int] = None
 
-    def __init__(self, storage):
+    def __init__(self, storage: BlobStore):
         self.storage = storage
 
-    def read_range(
-        self, uri: str, kbid: str, start: int, end: int
-    ) -> AsyncIterator[bytes]:
-        raise NotImplementedError()
-
     def iter_data(
         self, uri: str, kbid: str, headers: Optional[dict[str, str]] = None
     ) -> AsyncIterator[bytes]:
@@ -81,48 +69,6 @@ class FileStorageManager:
     async def delete_upload(self, uri, kbid):
         raise NotImplementedError()
 
-    async def full_download(self, content_length, content_type, upload_id):
-        return StreamingResponse(
-            self.iter_data(upload_id),
-            media_type=content_type,
-            headers={
-                "Content-Length": str(content_length),
-                "Content-Type": content_type,
-            },
-        )
-
-    async def range_download(
-        self, content_length, content_type, upload_id, range_header
-    ):
-        try:
-            start, _, end = range_header.split("bytes=")[-1].partition("-")
-            start = int(start)
-            if len(end) == 0:
-                # bytes=0- is valid
-                end = content_length - 1
-            end = int(end) + 1  # python is inclusive, http is exclusive
-        except (IndexError, ValueError):
-            # range errors fallback to full download
-            raise HTTPRangeNotSatisfiable(detail=f"Range not parsable {range_header}")
-        if start > end or start < 0:
-            raise HTTPRangeNotSatisfiable(detail="Invalid range {start}-{end}")
-        if end > content_length:
-            raise HTTPRangeNotSatisfiable(
-                detail="Invalid range {start}-{end}, too large end value"
-            )
-
-        logger.debug(f"Range request: {range_header}")
-        headers = {
-            "Content-Range": f"bytes {start}-{end - 1}/{content_length}",
-            "Content-Type": content_type,
-        }
-
-        return StreamingResponse(
-            self.read_range(upload_id, start, end),
-            media_type=content_type,
-            headers=headers,
-        )
-
     async def iterate_body_chunks(self, request, chunk_size):
         partial = b""
         remaining = b""
@@ -146,3 +92,6 @@ class FileStorageManager:
 
         if partial or remaining:
             yield partial + remaining
+
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        raise NotImplementedError()

nucliadb/writer/vectorsets.py

@@ -0,0 +1,125 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 nucliadb import learning_proxy
+from nucliadb.common import datamanagers
+from nucliadb.ingest.orm.exceptions import VectorSetConflict
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
+from nucliadb.writer import logger
+from nucliadb_protos import knowledgebox_pb2
+from nucliadb_telemetry import errors
+from nucliadb_utils.utilities import get_storage
+
+
+async def add(kbid: str, vectorset_id: str) -> None:
+    # First off, add the vectorset to the learning configuration if it's not already there
+    lconfig = await learning_proxy.get_configuration(kbid)
+    assert lconfig is not None
+    semantic_models = lconfig.model_dump()["semantic_models"]
+    if vectorset_id not in semantic_models:
+        semantic_models.append(vectorset_id)
+        await learning_proxy.update_configuration(kbid, {"semantic_models": semantic_models})
+        lconfig = await learning_proxy.get_configuration(kbid)
+        assert lconfig is not None
+
+    # Then, add the vectorset to the index if it's not already there
+    async with datamanagers.with_rw_transaction() as txn:
+        kbobj = KnowledgeBox(txn, await get_storage(), kbid)
+        vectorset_config = get_vectorset_config(lconfig, vectorset_id)
+        try:
+            await kbobj.create_vectorset(vectorset_config)
+            await txn.commit()
+        except VectorSetConflict:
+            # Vectorset already exists, nothing to do
+            return
+
+
+async def delete(kbid: str, vectorset_id: str) -> None:
+    lconfig = await learning_proxy.get_configuration(kbid)
+    if lconfig is not None:
+        semantic_models = lconfig.model_dump()["semantic_models"]
+        if vectorset_id in semantic_models:
+            semantic_models.remove(vectorset_id)
+            await learning_proxy.update_configuration(kbid, {"semantic_models": semantic_models})
+    try:
+        async with datamanagers.with_rw_transaction() as txn:
+            kbobj = KnowledgeBox(txn, await get_storage(), kbid)
+            await kbobj.delete_vectorset(vectorset_id=vectorset_id)
+            await txn.commit()
+    except Exception as ex:
+        errors.capture_exception(ex)
+        logger.exception(
+            "Could not delete vectorset from index", extra={"kbid": kbid, "vectorset_id": vectorset_id}
+        )
+
+
+def get_vectorset_config(
+    learning_config: learning_proxy.LearningConfiguration, vectorset_id: str
+) -> knowledgebox_pb2.VectorSetConfig:
+    """
+    Create a VectorSetConfig from a LearningConfiguration for a given vectorset_id
+    """
+    vectorset_config = knowledgebox_pb2.VectorSetConfig(vectorset_id=vectorset_id)
+    vectorset_index_config = knowledgebox_pb2.VectorIndexConfig(
+        vector_type=knowledgebox_pb2.VectorType.DENSE_F32,
+    )
+    model_config = learning_config.semantic_model_configs[vectorset_id]
+
+    # Parse similarity function
+    parsed_similarity = learning_proxy.SimilarityFunction(model_config.similarity)
+    if parsed_similarity == learning_proxy.SimilarityFunction.COSINE.value:
+        vectorset_index_config.similarity = knowledgebox_pb2.VectorSimilarity.COSINE
+    elif parsed_similarity == learning_proxy.SimilarityFunction.DOT.value:
+        vectorset_index_config.similarity = knowledgebox_pb2.VectorSimilarity.DOT
+    else:
+        raise ValueError(
+            f"Unknown similarity function {model_config.similarity}, parsed as {parsed_similarity}"
+        )
+
+    # Parse vector dimension
+    vectorset_index_config.vector_dimension = model_config.size
+
+    # Parse matryoshka dimensions
+    if len(model_config.matryoshka_dims) > 0:
+        vectorset_index_config.normalize_vectors = True
+        vectorset_config.matryoshka_dimensions.extend(model_config.matryoshka_dims)
+    else:
+        vectorset_index_config.normalize_vectors = False
+    vectorset_config.vectorset_index_config.CopyFrom(vectorset_index_config)
+    return vectorset_config

nucliadb-6.2.1.post2777.dist-info/METADATA

@@ -0,0 +1,148 @@
+Metadata-Version: 2.2
+Name: nucliadb
+Version: 6.2.1.post2777
+Home-page: https://docs.nuclia.dev/docs/management/nucliadb/intro
+Author: NucliaDB Community
+Author-email: nucliadb@nuclia.com
+License: BSD
+Project-URL: Nuclia, https://nuclia.com
+Project-URL: Github, https://github.com/nuclia/nucliadb
+Project-URL: Slack, https://nuclia-community.slack.com
+Project-URL: API Reference, https://docs.nuclia.dev/docs/api
+Keywords: search,semantic,AI
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3 :: Only
+Requires-Python: >=3.9, <4
+Description-Content-Type: text/markdown
+Requires-Dist: nucliadb-telemetry[all]>=6.2.1.post2777
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post2777
+Requires-Dist: nucliadb-protos>=6.2.1.post2777
+Requires-Dist: nucliadb-models>=6.2.1.post2777
+Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
+Requires-Dist: nucliadb-node-binding>=2.26.0
+Requires-Dist: nuclia-models>=0.24.2
+Requires-Dist: uvicorn
+Requires-Dist: argdantic
+Requires-Dist: aiohttp>=3.11.11
+Requires-Dist: lru-dict>=1.1.7
+Requires-Dist: backoff
+Requires-Dist: aiofiles>=0.8.0
+Requires-Dist: psutil>=5.9.7
+Requires-Dist: types-psutil>=5.9.5.17
+Requires-Dist: types-aiofiles>=0.8.3
+Requires-Dist: protobuf>=4.22.3
+Requires-Dist: types-protobuf<5,>=4.24
+Requires-Dist: grpcio<1.63.0,>=1.44.0
+Requires-Dist: grpcio-health-checking<1.63.0,>=1.44.0
+Requires-Dist: grpcio-channelz<1.63.0,>=1.44.0
+Requires-Dist: grpcio-status<1.63.0,>=1.44.0
+Requires-Dist: grpcio-tools<1.63.0,>=1.44.0
+Requires-Dist: grpcio-testing<1.63.0,>=1.44.0
+Requires-Dist: grpcio-reflection<1.63.0,>=1.44.0
+Requires-Dist: orjson>=3.6.7
+Requires-Dist: types-setuptools
+Requires-Dist: pydantic>=2.6
+Requires-Dist: pydantic-settings>=2.2
+Requires-Dist: aiobotocore>=2.9.0
+Requires-Dist: botocore>=1.34.0
+Requires-Dist: google-cloud-storage
+Requires-Dist: gcloud
+Requires-Dist: oauth2client
+Requires-Dist: jwcrypto>=1.5.6
+Requires-Dist: pyyaml>=5.1
+Requires-Dist: fastapi-versioning>=0.10.0
+Requires-Dist: fastapi>=0.95.2
+Requires-Dist: sentry-sdk>=2.8.0
+Requires-Dist: pyjwt>=2.4.0
+Requires-Dist: mmh3>=3.0.0
+Requires-Dist: httpx>=0.23.0
+Requires-Dist: grpc-stubs>=1.44.0
+Requires-Dist: aiodns>=3.0.0
+Requires-Dist: types-orjson
+Requires-Dist: psycopg[binary,pool]
+Requires-Dist: multidict>=6.0.4
+Requires-Dist: deprecated>=1.2.12
+Requires-Dist: asgiref>=3.3.2
+Requires-Dist: jmespath>=1.0.0
+Requires-Dist: idna>=3.3
+Requires-Dist: sniffio>=1.2.0
+Requires-Dist: async_lru>=2.0.4
+Requires-Dist: async-timeout>=4.0.3
+Requires-Dist: cachetools>=5.3.2
+Requires-Dist: types-cachetools>=5.3.0.5
+Requires-Dist: kubernetes_asyncio<30.0.0
+Provides-Extra: redis
+Requires-Dist: redis>=4.3.4; extra == "redis"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: project-url
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+
+# nucliadb
+
+This module contains most of the Python components for NucliaDB:
+
+- ingest
+- reader
+- writer
+- search
+- train
+
+# NucliaDB Migrations
+
+This module is used to manage NucliaDB Migrations.
+
+All migrations will be provided in the `migrations` folder and have a filename
+that follows the structure: `[sequence]_[migration name].py`.
+Where `sequence` is the order the migration should be run in with zero padding.
+Example: `0001_migrate_data.py`.
+
+Each migration should have the following:
+
+```python
+from nucliadb.migrator.context import ExecutionContext
+
+
+async def migrate(context: ExecutionContext) -> None:
+    """
+    Non-kb type of migration. Migrate global data.
+    """
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    """
+    Migrate kb.
+
+    Must have both types of migrations.
+    """
+```
+
+
+## How migrations are managed
+
+- All migrations utilize a distributed lock to prevent simulateously running jobs
+- Global migration state:
+    - current version
+    - target version
+    - KBs to migrate
+- KB Migration State:
+    - current version
+
+- Migrations are currently run with a deployment and will be continuously retried on failure.
+- Running migrations in a deployment is to make sure a migration does not prevent code deployment.