nucliadb-utils 6.9.0.post5076__py3-none-any.whl → 6.9.2.post5282__py3-none-any.whl

nucliadb_utils/const.py

@@ -41,4 +41,4 @@ class Features:
     SKIP_EXTERNAL_INDEX = "nucliadb_skip_external_index"
     LOG_REQUEST_PAYLOADS = "nucliadb_log_request_payloads"
     IGNORE_EXTRACTED_IN_SEARCH = "nucliadb_ignore_extracted_in_search"
-    DEBUG_MISSING_VECTORS = "nucliadb_debug_missing_vectors"
+    REBALANCE_ENABLED = "nucliadb_rebalance"

nucliadb_utils/featureflagging.py

@@ -45,10 +45,7 @@ DEFAULT_FLAG_DATA: dict[str, Any] = {
         "rollout": 0,
         "variants": {"environment": ["local"]},
     },
-    const.Features.DEBUG_MISSING_VECTORS: {
-        "rollout": 0,
-        "variants": {"environment": ["local"]},
-    },
+    const.Features.REBALANCE_ENABLED: {"rollout": 0, "variants": {"environment": ["local"]}},
 }
 
 

nucliadb_utils/tests/fixtures.py

@@ -23,12 +23,14 @@ from typing import Any, Iterator, Type
 from unittest.mock import Mock
 
 import pytest
+from pytest import FixtureRequest
 from pytest_lazy_fixtures import lazy_fixture
 
 from nucliadb_utils.storages.azure import AzureStorage
 from nucliadb_utils.storages.gcs import GCSStorage
 from nucliadb_utils.storages.local import LocalStorage
 from nucliadb_utils.storages.s3 import S3Storage
+from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import Utility, clean_utility, set_utility
 
 
@@ -52,7 +54,7 @@ def hosted_nucliadb():
     nuclia_settings.onprem = original
 
 
-def get_testing_storage_backend():
+def get_testing_storage_backend() -> str:
     """
     Default to gcs for linux users and s3 for macOS users. This is because some
     tests fail on macOS with the gcs backend with a weird nidx error (to be looked into).
@@ -72,7 +74,7 @@ def lazy_storage_fixture():
 
 
 @pytest.fixture(scope="function", params=lazy_storage_fixture())
-async def storage(request):
+def storage(request: FixtureRequest) -> Iterator[Storage]:
     """
     Generic storage fixture that allows us to run the same tests for different storage backends.
     """

nucliadb_utils/tests/nats.py

@@ -18,12 +18,16 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import socket
+from typing import AsyncIterator, Iterator
 
 import nats
 import pytest
 from pytest_docker_fixtures import images  # type: ignore
 from pytest_docker_fixtures.containers._base import BaseImage  # type: ignore
 
+from nucliadb_utils.nats import NatsConnectionManager
+from nucliadb_utils.utilities import start_nats_manager, stop_nats_manager
+
 images.settings["nats"] = {
     "image": "nats",
     "version": "2.10.21",
@@ -48,7 +52,7 @@ nats_image = NatsImage()
 
 
 @pytest.fixture(scope="session")
-def natsd():  # pragma: no cover
+def natsd() -> Iterator[str]:
     nats_host, nats_port = nats_image.run()
     print("Started natsd docker")
     yield f"nats://{nats_host}:{nats_port}"
@@ -56,10 +60,17 @@ def natsd():  # pragma: no cover
 
 
 @pytest.fixture(scope="function")
-async def nats_server(natsd: str):
+async def nats_server(natsd: str) -> AsyncIterator[str]:
     yield natsd
 
     # cleanup nats
     nc = await nats.connect(servers=[natsd])
     await nc.drain()
     await nc.close()
+
+
+@pytest.fixture(scope="function")
+async def nats_manager(nats_server: str) -> AsyncIterator[NatsConnectionManager]:
+    ncm = await start_nats_manager("nucliadb_tests", [nats_server], None)
+    yield ncm
+    await stop_nats_manager()

nucliadb_utils/utilities.py

@@ -29,7 +29,6 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union, cast
 from nucliadb_protos.writer_pb2_grpc import WriterStub
 from nucliadb_telemetry.metrics import Counter
 from nucliadb_utils import featureflagging
-from nucliadb_utils.aiopynecone.client import PineconeSession
 from nucliadb_utils.audit.audit import AuditStorage
 from nucliadb_utils.audit.basic import BasicAuditStorage
 from nucliadb_utils.audit.stream import StreamAuditStorage
@@ -85,7 +84,6 @@ class Utility(str, Enum):
     MAINDB_DRIVER = "driver"
     USAGE = "usage"
     ENDECRYPTOR = "endecryptor"
-    PINECONE_SESSION = "pinecone_session"
     NIDX = "nidx"
 
 
@@ -367,6 +365,11 @@ async def stop_audit_utility():
 async def start_nats_manager(
     service_name: str, nats_servers: list[str], nats_creds: Optional[str] = None
 ) -> NatsConnectionManager:
+    util = get_utility(Utility.NATS_MANAGER)
+    if util is not None:
+        logger.warning("Warning, nats manager utility was already set, ignoring")
+        return util
+
     nats_manager = NatsConnectionManager(
         service_name=service_name,
         nats_servers=nats_servers,
@@ -439,16 +442,3 @@ def get_endecryptor() -> EndecryptorUtility:
         ) from ex
     set_utility(Utility.ENDECRYPTOR, util)
     return util
-
-
-def get_pinecone() -> PineconeSession:
-    util = get_utility(Utility.PINECONE_SESSION)
-    if util is not None:
-        return util
-    util = PineconeSession()
-    set_utility(Utility.PINECONE_SESSION, util)
-    return util
-
-
-def clean_pinecone():
-    clean_utility(Utility.PINECONE_SESSION)

{nucliadb_utils-6.9.0.post5076.dist-info → nucliadb_utils-6.9.2.post5282.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb_utils
-Version: 6.9.0.post5076
+Version: 6.9.2.post5282
 Summary: NucliaDB util library
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -8,13 +8,12 @@ Project-URL: Homepage, https://nuclia.com
 Project-URL: Repository, https://github.com/nuclia/nucliadb
 Classifier: Development Status :: 4 - Beta
 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
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: <4,>=3.9
+Requires-Python: <4,>=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: pydantic>=2.6
 Requires-Dist: pydantic-settings>=2.2
@@ -27,8 +26,8 @@ Requires-Dist: nats-py[nkeys]>=2.6.0
 Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=6.9.0.post5076
-Requires-Dist: nucliadb-telemetry>=6.9.0.post5076
+Requires-Dist: nucliadb-protos>=6.9.2.post5282
+Requires-Dist: nucliadb-telemetry>=6.9.2.post5282
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.9.0.post5076.dist-info → nucliadb_utils-6.9.2.post5282.dist-info}/RECORD

@@ -1,10 +1,10 @@
 nucliadb_utils/__init__.py,sha256=EvBCH1iTODe-AgXm48aj4kVUt_Std3PeL8QnwimR5wI,895
 nucliadb_utils/asyncio_utils.py,sha256=h8Y-xpcFFRgNzaiIW0eidz7griAQa7ggbNk34-tAt2c,2888
 nucliadb_utils/authentication.py,sha256=5_b323v2ylJaJvM_0coeSQEtnD-p9IGD-6CPA6IXhik,6471
-nucliadb_utils/const.py,sha256=kgTPXlvVxJB_HPz9GkEigE94UNEkrnMECgrqhPW2Qi4,1537
+nucliadb_utils/const.py,sha256=YtWadXGm044MbwLcfEImNC6skh3e4LKZDu6hjlO0qMU,1521
 nucliadb_utils/debug.py,sha256=Q56Nx9Dp7V2ae3CU2H0ztaZcHTJXdlflPLKLeOPZ170,2436
 nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,1094
-nucliadb_utils/featureflagging.py,sha256=kLIJT-AJe8AqTuL0sY9nn9r5Soax8EvqPangM6UygfM,2380
+nucliadb_utils/featureflagging.py,sha256=ctd9Nqm_nhoedMIV2GC819-cSP5GlkLYXCRE0DbwxYU,2353
 nucliadb_utils/grpc.py,sha256=apu0uePnkGHCAT7GRQ9YZfRYyFj26kJ440i8jitbM3U,3314
 nucliadb_utils/helpers.py,sha256=eed7_E1MKh9eW3CpqOXka3OvLw5C9eJGC_R-1MPYdfY,3336
 nucliadb_utils/nats.py,sha256=U21Cfg36_IHd3ZLXEC4eZ7nZ1Soh_ZNFFwjryNyd2-8,15248
@@ -15,11 +15,7 @@ nucliadb_utils/settings.py,sha256=H9yKrHPR5emTxai-D4owg4CjE4_-E0qR0HyuHERQNH4,84
 nucliadb_utils/signals.py,sha256=lo_Mk12NIX5Au--3H3WObvDOXq_OMurql2qiC2TnAao,2676
 nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=l3ZvrITYMnAs_fv1OOC-1nDZxWPG5qmbBhzvuC3DUzQ,8039
-nucliadb_utils/utilities.py,sha256=Q5sJmoX7JSnAId1BC9Lr3H9IQ85YmFwUMtw9Lco_brQ,16044
-nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/aiopynecone/client.py,sha256=MPyHnDXwhukJr7U3CJh7BpsekfSuOkyM4g5b9LLtzc8,22941
-nucliadb_utils/aiopynecone/exceptions.py,sha256=fUErx3ceKQK1MUbOnYcZhIzpNe8UVAptZE9JIRDLXDE,4000
-nucliadb_utils/aiopynecone/models.py,sha256=XkNIZx4bxdbVo9zYVn8IRp70q4DWUMWN79ybGloFj2Q,3492
+nucliadb_utils/utilities.py,sha256=SjPnCwCUH_lWUKSOZQp9vIcTYmLP0yL_UC8nPwlnds4,15817
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=xmJJiAGG8rPGADwD9gXN9-QJ80GeGvqmY-kCwEf6PiQ,3598
 nucliadb_utils/audit/basic.py,sha256=fcCYvoFSGVbbB8cSCnm95bN2rf1AAeuWhGfh5no0S-Y,4246
@@ -52,12 +48,12 @@ nucliadb_utils/storages/utils.py,sha256=F4Iboa_0_bhDQr-JOKD9sGPld_-hKwJW5ptyZdn9
 nucliadb_utils/tests/__init__.py,sha256=Oo9CAE7B0eW5VHn8sHd6o30SQzOWUhktLPRXdlDOleA,1456
 nucliadb_utils/tests/asyncbenchmark.py,sha256=vrX_x9ifCXi18PfNShc23w9x_VUiB_Ph-2nuolh9z3Q,10707
 nucliadb_utils/tests/azure.py,sha256=rt1KRSYZW1EYhKy4Q0i7IEL9vdoOU6BYw2__S51YfGg,5039
-nucliadb_utils/tests/fixtures.py,sha256=4lzz-khYvbGzdbT18IG6KKg40f7CVex2q3ho88I-jL8,3799
+nucliadb_utils/tests/fixtures.py,sha256=-OeR4NhtXveBqN6sZa7KVaNwyUdsvJUxS_yFmIgF148,3923
 nucliadb_utils/tests/gcs.py,sha256=JNqp5ymeNNU9Ci8rNYTh7-VqP4fjybElhyB3ap7EV1c,4721
 nucliadb_utils/tests/local.py,sha256=z9E11_ol1mu7N8Y6PkjKl-WMPPMl7JqQbDj3uhVa1A0,1933
-nucliadb_utils/tests/nats.py,sha256=RWHjwqq5esuO7OFbP24yYX1cXnpPLcWJwDUdmwCpH28,1897
+nucliadb_utils/tests/nats.py,sha256=rbTaC6kv-u6SdZ7N-XBEGS40XCRiUmFUsKHIYWJfxTs,2325
 nucliadb_utils/tests/s3.py,sha256=kz9ULxrAYLVslZ59I8dtweZ9DJz5R8Ioy2XYrveZzHw,3829
-nucliadb_utils-6.9.0.post5076.dist-info/METADATA,sha256=yZKmoZRlowkHvQl8q-eAlAVMF0mU17eGJe8Ql-kBnd8,2180
-nucliadb_utils-6.9.0.post5076.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb_utils-6.9.0.post5076.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.9.0.post5076.dist-info/RECORD,,
+nucliadb_utils-6.9.2.post5282.dist-info/METADATA,sha256=q0qd9MBC9AOpmBgjqMw2xwPB5z6Fy2P_sIEqH8wfRWk,2131
+nucliadb_utils-6.9.2.post5282.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb_utils-6.9.2.post5282.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.9.2.post5282.dist-info/RECORD,,

nucliadb_utils/aiopynecone/__init__.py

@@ -1,19 +0,0 @@
-# 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/>.
-#

nucliadb_utils/aiopynecone/client.py

@@ -1,636 +0,0 @@
-# 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/>.
-#
-import asyncio
-import json
-import logging
-import random
-from collections.abc import AsyncIterable, Iterable
-from itertools import islice
-from typing import Any, AsyncGenerator, Optional
-
-import backoff
-import httpx
-
-from nucliadb_telemetry.metrics import INF, Histogram, Observer
-from nucliadb_utils.aiopynecone.exceptions import (
-    PineconeAPIError,
-    PineconeRateLimitError,
-    RetriablePineconeAPIError,
-    raise_for_status,
-)
-from nucliadb_utils.aiopynecone.models import (
-    CreateIndexRequest,
-    CreateIndexResponse,
-    IndexDescription,
-    IndexStats,
-    ListResponse,
-    QueryResponse,
-    UpsertRequest,
-    Vector,
-)
-
-logger = logging.getLogger(__name__)
-
-upsert_batch_size_histogram = Histogram(
-    "pinecone_upsert_batch_size",
-    buckets=[10.0, 100.0, 200.0, 500.0, 1000.0, 5000.0, INF],
-)
-upsert_batch_count_histogram = Histogram(
-    "pinecone_upsert_batch_count",
-    buckets=[0.0, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 20.0, 30.0, 50.0, INF],
-)
-
-delete_batch_size_histogram = Histogram(
-    "pinecone_delete_batch_size",
-    buckets=[1.0, 5.0, 10.0, 20.0, 50.0, 100.0, 150.0, INF],
-)
-
-delete_batch_count_histogram = Histogram(
-    "pinecone_delete_batch_count",
-    buckets=[0.0, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 20.0, 30.0, 50.0, INF],
-)
-
-
-pinecone_observer = Observer(
-    "pinecone_client",
-    labels={"type": ""},
-    error_mappings={
-        "rate_limit": PineconeRateLimitError,
-    },
-)
-
-DEFAULT_TIMEOUT = 30
-CONTROL_PLANE_BASE_URL = "https://api.pinecone.io/"
-BASE_API_HEADERS = {
-    "Content-Type": "application/json",
-    "Accept": "application/json",
-    # This is needed so that it is easier for Pinecone to track which api requests
-    # are coming from the Nuclia integration:
-    # https://docs.pinecone.io/integrations/build-integration/attribute-usage-to-your-integration
-    "User-Agent": "source_tag=nuclia",
-    "X-Pinecone-API-Version": "2024-10",
-}
-MEGA_BYTE = 1024 * 1024
-MAX_UPSERT_PAYLOAD_SIZE = 2 * MEGA_BYTE
-MAX_DELETE_BATCH_SIZE = 1000
-MAX_LIST_PAGE_SIZE = 100
-
-
-RETRIABLE_EXCEPTIONS = (
-    PineconeRateLimitError,
-    RetriablePineconeAPIError,
-    httpx.ConnectError,
-    httpx.NetworkError,
-    httpx.WriteTimeout,
-    httpx.ReadTimeout,
-)
-
-
-class ControlPlane:
-    """
-    Client for interacting with the Pinecone control plane API.
-    https://docs.pinecone.io/reference/api/control-plane
-    """
-
-    def __init__(self, api_key: str, http_session: httpx.AsyncClient):
-        self.api_key = api_key
-        self.http_session = http_session
-
-    @pinecone_observer.wrap({"type": "create_index"})
-    async def create_index(
-        self,
-        name: str,
-        dimension: int,
-        metric: str = "dotproduct",
-        serverless_cloud: Optional[dict[str, str]] = None,
-    ) -> str:
-        """
-        Create a new index in Pinecone. It can only create serverless indexes on the AWS us-east-1 region.
-        Params:
-        - `name`: The name of the index.
-        - `dimension`: The dimension of the vectors in the index.
-        - `metric`: The similarity metric to use. Default is "dotproduct".
-        - `serverless_cloud`: The serverless provider to use. Default is AWS us-east-1.
-        Returns:
-        - The index host to be used for data plane operations.
-        """
-        serverless_cloud = serverless_cloud or {"cloud": "aws", "region": "us-east-1"}
-        payload = CreateIndexRequest(
-            name=name,
-            dimension=dimension,
-            metric=metric,
-            spec={"serverless": serverless_cloud},
-        )
-        headers = {"Api-Key": self.api_key}
-        http_response = await self.http_session.post(
-            "/indexes", json=payload.model_dump(), headers=headers
-        )
-        raise_for_status("create_index", http_response)
-        response = CreateIndexResponse.model_validate(http_response.json())
-        return response.host
-
-    @pinecone_observer.wrap({"type": "delete_index"})
-    async def delete_index(self, name: str) -> None:
-        """
-        Delete an index in Pinecone.
-        Params:
-        - `name`: The name of the index to delete.
-        """
-        headers = {"Api-Key": self.api_key}
-        response = await self.http_session.delete(f"/indexes/{name}", headers=headers)
-        if response.status_code == 404:  # pragma: no cover
-            logger.warning("Pinecone index not found.", extra={"index_name": name})
-            return
-        raise_for_status("delete_index", response)
-
-    @pinecone_observer.wrap({"type": "describe_index"})
-    async def describe_index(self, name: str) -> IndexDescription:
-        """
-        Describe an index in Pinecone.
-        Params:
-        - `name`: The name of the index to describe.
-        """
-        headers = {"Api-Key": self.api_key}
-        response = await self.http_session.get(f"/indexes/{name}", headers=headers)
-        raise_for_status("describe_index", response)
-        return IndexDescription.model_validate(response.json())
-
-
-class DataPlane:
-    """
-    Client for interacting with the Pinecone data plane API, hosted by an index host.
-    https://docs.pinecone.io/reference/api/data-plane
-    """
-
-    def __init__(
-        self, api_key: str, index_host_session: httpx.AsyncClient, timeout: Optional[float] = None
-    ):
-        """
-        Params:
-        - `api_key`: The Pinecone API key.
-        - `index_host_session`: The http session for the index host.
-        - `timeout`: The default timeout for all requests. If not set, the default timeout from httpx.AsyncClient is used.
-        """
-        self.api_key = api_key
-        self.http_session = index_host_session
-        self.client_timeout = timeout
-        self._upsert_batch_size: Optional[int] = None
-
-    def _get_request_timeout(self, timeout: Optional[float] = None) -> Optional[float]:
-        return timeout or self.client_timeout
-
-    @pinecone_observer.wrap({"type": "stats"})
-    async def stats(self, filter: Optional[dict[str, Any]] = None) -> IndexStats:
-        """
-        Get the index stats.
-        Params:
-        - `filter`: to filter the stats by their metadata. See:
-        https://docs.pinecone.io/reference/api/2024-07/data-plane/describeindexstats
-        """
-        post_kwargs: dict[str, Any] = {
-            "headers": {"Api-Key": self.api_key},
-        }
-        if filter is not None:
-            post_kwargs["json"] = {
-                "filter": filter,
-            }
-        response = await self.http_session.post("/describe_index_stats", **post_kwargs)
-        raise_for_status("stats", response)
-        return IndexStats.model_validate(response.json())
-
-    @backoff.on_exception(
-        backoff.expo,
-        RETRIABLE_EXCEPTIONS,
-        jitter=backoff.random_jitter,
-        max_tries=4,
-    )
-    @pinecone_observer.wrap({"type": "upsert"})
-    async def upsert(self, vectors: list[Vector], timeout: Optional[float] = None) -> None:
-        """
-        Upsert vectors into the index.
-        Params:
-        - `vectors`: The vectors to upsert.
-        - `timeout`: to control the request timeout. If not set, the default timeout is used.
-        """
-        if len(vectors) == 0:
-            # Nothing to upsert.
-            return
-        upsert_batch_size_histogram.observe(len(vectors))
-        headers = {"Api-Key": self.api_key}
-        payload = UpsertRequest(vectors=vectors)
-        post_kwargs: dict[str, Any] = {
-            "headers": headers,
-            "json": payload.model_dump(),
-        }
-        request_timeout = self._get_request_timeout(timeout)
-        if request_timeout is not None:
-            post_kwargs["timeout"] = timeout
-        response = await self.http_session.post("/vectors/upsert", **post_kwargs)
-        raise_for_status("upsert", response)
-
-    def _estimate_upsert_batch_size(self, vectors: list[Vector]) -> int:
-        """
-        Estimate a batch size so that the upsert payload does not exceed the hard limit.
-        https://docs.pinecone.io/reference/quotas-and-limits#hard-limits
-        """
-        if self._upsert_batch_size is not None:
-            # Return the cached value.
-            return self._upsert_batch_size
-        # Take the dimension of the first vector as the vector dimension.
-        # Assumes all vectors have the same dimension.
-        vector_dimension = len(vectors[0].values)
-        # Estimate the metadata size by taking the average of 20 random vectors.
-        metadata_sizes = []
-        for _ in range(20):
-            metadata_sizes.append(len(json.dumps(random.choice(vectors).metadata)))
-        average_metadata_size = sum(metadata_sizes) / len(metadata_sizes)
-        # Estimate the size of the vector payload. 4 bytes per float.
-        vector_size = 4 * vector_dimension + average_metadata_size
-        # Cache the value.
-        self._upsert_batch_size = max(int(MAX_UPSERT_PAYLOAD_SIZE // vector_size), 1)
-        return self._upsert_batch_size
-
-    @pinecone_observer.wrap({"type": "upsert_in_batches"})
-    async def upsert_in_batches(
-        self,
-        vectors: list[Vector],
-        batch_size: Optional[int] = None,
-        max_parallel_batches: int = 1,
-        batch_timeout: Optional[float] = None,
-    ) -> None:
-        """
-        Upsert vectors in batches.
-        Params:
-        - `vectors`: The vectors to upsert.
-        - `batch_size`: to control the number of vectors in each batch.
-        - `max_parallel_batches`: to control the number of batches sent concurrently.
-        - `batch_timeout`: to control the request timeout for each batch.
-        """
-        if batch_size is None:
-            batch_size = self._estimate_upsert_batch_size(vectors)
-
-        semaphore = asyncio.Semaphore(max_parallel_batches)
-
-        async def _upsert_batch(batch):
-            async with semaphore:
-                await self.upsert(vectors=batch, timeout=batch_timeout)
-
-        tasks = []
-        for batch in batchify(vectors, batch_size):
-            tasks.append(asyncio.create_task(_upsert_batch(batch)))
-
-        upsert_batch_count_histogram.observe(len(tasks))
-
-        if len(tasks) > 0:
-            await asyncio.gather(*tasks)
-
-    @backoff.on_exception(
-        backoff.expo,
-        RETRIABLE_EXCEPTIONS,
-        jitter=backoff.random_jitter,
-        max_tries=4,
-    )
-    @pinecone_observer.wrap({"type": "delete"})
-    async def delete(self, ids: list[str], timeout: Optional[float] = None) -> None:
-        """
-        Delete vectors by their ids.
-        Maximum number of ids in a single request is 1000.
-        Params:
-        - `ids`: The ids of the vectors to delete.
-        - `timeout`: to control the request timeout. If not set, the default timeout is used.
-        """
-        if len(ids) > MAX_DELETE_BATCH_SIZE:
-            raise ValueError(f"Maximum number of ids in a single request is {MAX_DELETE_BATCH_SIZE}.")
-        if len(ids) == 0:  # pragma: no cover
-            return
-
-        delete_batch_size_histogram.observe(len(ids))
-        headers = {"Api-Key": self.api_key}
-
-        # This is a temporary log info to hunt down a bug.
-        rids = {vid.split("/")[0] for vid in ids}
-        logger.info(f"Deleting vectors from resources: {list(rids)}")
-
-        payload = {"ids": ids}
-        post_kwargs: dict[str, Any] = {
-            "headers": headers,
-            "json": payload,
-        }
-        request_timeout = self._get_request_timeout(timeout)
-        if request_timeout is not None:
-            post_kwargs["timeout"] = timeout
-        response = await self.http_session.post("/vectors/delete", **post_kwargs)
-        raise_for_status("delete", response)
-
-    @backoff.on_exception(
-        backoff.expo,
-        RETRIABLE_EXCEPTIONS,
-        jitter=backoff.random_jitter,
-        max_tries=4,
-    )
-    @pinecone_observer.wrap({"type": "list_page"})
-    async def list_page(
-        self,
-        id_prefix: Optional[str] = None,
-        limit: int = MAX_LIST_PAGE_SIZE,
-        pagination_token: Optional[str] = None,
-        timeout: Optional[float] = None,
-    ) -> ListResponse:
-        """
-        List vectors in a paginated manner.
-        Params:
-        - `id_prefix`: to filter vectors by their id prefix.
-        - `limit`: to control the number of vectors fetched in each page.
-        - `pagination_token`: to fetch the next page. The token is provided in the response
-           if there are more pages to fetch.
-        - `timeout`: to control the request timeout. If not set, the default timeout is used.
-        """
-        if limit > MAX_LIST_PAGE_SIZE:  # pragma: no cover
-            raise ValueError(f"Maximum limit is {MAX_LIST_PAGE_SIZE}.")
-        headers = {"Api-Key": self.api_key}
-        params = {"limit": str(limit)}
-        if id_prefix is not None:
-            params["prefix"] = id_prefix
-        if pagination_token is not None:
-            params["paginationToken"] = pagination_token
-
-        post_kwargs: dict[str, Any] = {
-            "headers": headers,
-            "params": params,
-        }
-        request_timeout = self._get_request_timeout(timeout)
-        if request_timeout is not None:
-            post_kwargs["timeout"] = timeout
-        response = await self.http_session.get(
-            "/vectors/list",
-            **post_kwargs,
-        )
-        raise_for_status("list_page", response)
-        return ListResponse.model_validate(response.json())
-
-    async def list_all(
-        self,
-        id_prefix: Optional[str] = None,
-        page_size: int = MAX_LIST_PAGE_SIZE,
-        page_timeout: Optional[float] = None,
-    ) -> AsyncGenerator[str, None]:
-        """
-        Iterate over all vector ids from the index in a paginated manner.
-        Params:
-        - `id_prefix`: to filter vectors by their id prefix.
-        - `page_size`: to control the number of vectors fetched in each page.
-        - `page_timeout`: to control the request timeout for each page. If not set, the default timeout is used.
-        """
-        pagination_token = None
-        while True:
-            response = await self.list_page(
-                id_prefix=id_prefix,
-                limit=page_size,
-                pagination_token=pagination_token,
-                timeout=page_timeout,
-            )
-            for vector_id in response.vectors:
-                yield vector_id.id
-            if response.pagination is None:
-                break
-            pagination_token = response.pagination.next
-
-    @backoff.on_exception(
-        backoff.expo,
-        RETRIABLE_EXCEPTIONS,
-        jitter=backoff.random_jitter,
-        max_tries=4,
-    )
-    @pinecone_observer.wrap({"type": "delete_all"})
-    async def delete_all(self, timeout: Optional[float] = None):
-        """
-        Delete all vectors in the index.
-        Params:
-        - `timeout`: to control the request timeout. If not set, the default timeout is used.
-        """
-        headers = {"Api-Key": self.api_key}
-        payload = {"deleteAll": True, "ids": [], "namespace": ""}
-        post_kwargs: dict[str, Any] = {
-            "headers": headers,
-            "json": payload,
-        }
-        request_timeout = self._get_request_timeout(timeout)
-        if request_timeout is not None:
-            post_kwargs["timeout"] = timeout
-        response = await self.http_session.post("/vectors/delete", **post_kwargs)
-        try:
-            raise_for_status("delete_all", response)
-        except PineconeAPIError as err:
-            if err.http_status_code == 404 and err.code == 5:  # pragma: no cover
-                # Namespace not found. No vectors to delete.
-                return
-            raise
-
-    @pinecone_observer.wrap({"type": "delete_by_id_prefix"})
-    async def delete_by_id_prefix(
-        self,
-        id_prefix: str,
-        batch_size: int = MAX_DELETE_BATCH_SIZE,
-        max_parallel_batches: int = 1,
-        batch_timeout: Optional[float] = None,
-    ) -> None:
-        """
-        Delete vectors by their id prefix. It lists all vectors with the given prefix and deletes them in batches.
-        Params:
-        - `id_prefix`: to filter vectors by their id prefix.
-        - `batch_size`: to control the number of vectors deleted in each batch. Maximum is 1000.
-        - `max_parallel_batches`: to control the number of batches sent concurrently.
-        - `batch_timeout`: to control the request timeout for each batch.
-        """
-        if batch_size > MAX_DELETE_BATCH_SIZE:
-            logger.warning(f"Batch size {batch_size} is too large. Limiting to {MAX_DELETE_BATCH_SIZE}.")
-            batch_size = MAX_DELETE_BATCH_SIZE
-
-        semaphore = asyncio.Semaphore(max_parallel_batches)
-
-        async def _delete_batch(batch):
-            async with semaphore:
-                await self.delete(ids=batch, timeout=batch_timeout)
-
-        tasks = []
-        async_iterable = self.list_all(id_prefix=id_prefix, page_timeout=batch_timeout)
-        async for batch in async_batchify(async_iterable, batch_size):
-            tasks.append(asyncio.create_task(_delete_batch(batch)))
-
-        delete_batch_count_histogram.observe(len(tasks))
-
-        if len(tasks) > 0:
-            await asyncio.gather(*tasks)
-
-    @backoff.on_exception(
-        backoff.expo,
-        RETRIABLE_EXCEPTIONS,
-        jitter=backoff.random_jitter,
-        max_tries=4,
-    )
-    @pinecone_observer.wrap({"type": "query"})
-    async def query(
-        self,
-        vector: list[float],
-        top_k: int = 20,
-        include_values: bool = False,
-        include_metadata: bool = False,
-        filter: Optional[dict[str, Any]] = None,
-        timeout: Optional[float] = None,
-    ) -> QueryResponse:
-        """
-        Query the index for similar vectors to the given vector.
-        Params:
-        - `vector`: The query vector.
-        - `top_k`: to control the number of similar vectors to return.
-        - `include_values`: to include the vector values in the response.
-        - `include_metadata`: to include the vector metadata in the response.
-        - `filter`: to filter the vectors by their metadata. See:
-           https://docs.pinecone.io/guides/data/filter-with-metadata#metadata-query-language
-        - `timeout`: to control the request timeout. If not set, the default timeout is used.
-        """
-        headers = {"Api-Key": self.api_key}
-        payload = {
-            "vector": vector,
-            "topK": top_k,
-            "includeValues": include_values,
-            "includeMetadata": include_metadata,
-        }
-        if filter:
-            payload["filter"] = filter
-        post_kwargs: dict[str, Any] = {
-            "headers": headers,
-            "json": payload,
-        }
-        request_timeout = self._get_request_timeout(timeout)
-        if request_timeout is not None:
-            post_kwargs["timeout"] = timeout
-        response = await self.http_session.post("/query", **post_kwargs)
-        raise_for_status("query", response)
-        return QueryResponse.model_validate(response.json())
-
-
-class PineconeSession:
-    """
-    Wrapper class that manages the sessions around all Pinecone http api interactions.
-    Holds a single control plane session and multiple data plane sessions, one for each index host.
-    """
-
-    def __init__(self):
-        self.control_plane_session = httpx.AsyncClient(
-            base_url=CONTROL_PLANE_BASE_URL, headers=BASE_API_HEADERS, timeout=DEFAULT_TIMEOUT
-        )
-        self.index_host_sessions = {}
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.finalize()
-
-    async def finalize(self):
-        if not self.control_plane_session.is_closed:
-            await self.control_plane_session.aclose()
-        for session in self.index_host_sessions.values():
-            if not session.is_closed:
-                await session.aclose()
-        self.index_host_sessions.clear()
-
-    def control_plane(self, api_key: str) -> ControlPlane:
-        return ControlPlane(api_key=api_key, http_session=self.control_plane_session)
-
-    def _get_index_host_session(self, index_host: str) -> httpx.AsyncClient:
-        """
-        Get a session for the given index host.
-        Cache http sessions so that they are reused for the same index host.
-        """
-        session = self.index_host_sessions.get(index_host, None)
-        if session is not None:
-            return session
-
-        base_url = index_host
-        if not index_host.startswith("https://"):
-            base_url = f"https://{index_host}/"
-
-        session = httpx.AsyncClient(
-            base_url=base_url,
-            headers=BASE_API_HEADERS,
-            timeout=DEFAULT_TIMEOUT,
-        )
-        self.index_host_sessions[index_host] = session
-        return session
-
-    def data_plane(self, api_key: str, index_host: str, timeout: Optional[float] = None) -> DataPlane:
-        index_host_session = self._get_index_host_session(index_host)
-        return DataPlane(api_key=api_key, index_host_session=index_host_session, timeout=timeout)
-
-
-def batchify(iterable: Iterable, batch_size: int):
-    """
-    Split an iterable into batches of batch_size
-    """
-    iterator = iter(iterable)
-    while True:
-        batch = list(islice(iterator, batch_size))
-        if not batch:
-            break
-        yield batch
-
-
-async def async_batchify(async_iterable: AsyncIterable, batch_size: int):
-    """
-    Split an async iterable into batches of batch_size
-    """
-    batch = []
-    async for item in async_iterable:
-        batch.append(item)
-        if len(batch) == batch_size:
-            yield batch
-            batch = []
-    if batch:
-        yield batch
-
-
-class FilterOperator:
-    """
-    Filter operators for metadata queries.
-    https://docs.pinecone.io/guides/data/filter-with-metadata#metadata-query-language
-    """
-
-    EQUALS = "$eq"
-    NOT_EQUALS = "$ne"
-    GREATER_THAN = "$gt"
-    GREATER_THAN_OR_EQUAL = "$gte"
-    LESS_THAN = "$lt"
-    LESS_THAN_OR_EQUAL = "$lte"
-    IN = "$in"
-    NOT_IN = "$nin"
-    EXISTS = "$exists"
-
-
-class LogicalOperator:
-    """
-    Logical operators for metadata queries.
-    https://docs.pinecone.io/guides/data/filter-with-metadata#metadata-query-language
-    """
-
-    AND = "$and"
-    OR = "$or"
-    NOT = "$not"

nucliadb_utils/aiopynecone/exceptions.py

@@ -1,131 +0,0 @@
-# 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 typing import Any, Optional
-
-import httpx
-
-from nucliadb_telemetry.metrics import Counter
-
-pinecone_errors_counter = Counter("pinecone_errors", labels={"type": "", "status_code": ""})
-
-
-class PineconeAPIError(Exception):
-    """
-    Generic Pinecone API error.
-    """
-
-    def __init__(
-        self,
-        http_status_code: int,
-        code: Optional[str] = None,
-        message: Optional[str] = None,
-        details: Optional[Any] = None,
-    ):
-        self.http_status_code = http_status_code
-        self.code = code or ""
-        self.message = message or ""
-        self.details = details or {}
-        exc_message = '[{http_status_code}] message="{message}" code={code} details={details}'.format(
-            http_status_code=http_status_code,
-            message=message,
-            code=code,
-            details=details,
-        )
-        super().__init__(exc_message)
-
-
-class RetriablePineconeAPIError(PineconeAPIError):
-    """
-    Raised when the client can retry the operation.
-    """
-
-    pass
-
-
-class PineconeRateLimitError(RetriablePineconeAPIError):
-    """
-    Raised when the client has exceeded the rate limit to be able to backoff and retry.
-    """
-
-    pass
-
-
-class PineconeNeedsPlanUpgradeError(PineconeAPIError):
-    """
-    Raised when the client needs to upgrade the plan to continue using the service.
-    """
-
-    pass
-
-
-class MetadataTooLargeError(ValueError):
-    """
-    Raised when the metadata of a vector to be upserted is too large.
-    """
-
-    pass
-
-
-def raise_for_status(operation: str, response: httpx.Response):
-    try:
-        response.raise_for_status()
-    except httpx.HTTPStatusError:
-        pinecone_errors_counter.inc(labels={"type": operation, "status_code": str(response.status_code)})
-        code = None
-        message = None
-        details = None
-        try:
-            resp_json = response.json()
-            error = resp_json.get("error") or {}
-            code = resp_json.get("code") or error.get("code")
-            message = resp_json.get("message") or error.get("message") or ""
-            details = resp_json.get("details") or error.get("details")
-        except Exception:  # pragma: no cover
-            message = response.text
-        if response.status_code == 429:
-            if "month" in message:
-                raise PineconeNeedsPlanUpgradeError(
-                    http_status_code=response.status_code,
-                    code=code,
-                    message=message,
-                    details=details,
-                )
-            raise PineconeRateLimitError(
-                http_status_code=response.status_code,
-                code=code,
-                message=message,
-                details=details,
-            )
-
-        if str(response.status_code).startswith("5"):
-            raise RetriablePineconeAPIError(
-                http_status_code=response.status_code,
-                code=code,
-                message=message,
-                details=details,
-            )
-
-        raise PineconeAPIError(
-            http_status_code=response.status_code,
-            code=code,
-            message=message,
-            details=details,
-        )

nucliadb_utils/aiopynecone/models.py

@@ -1,139 +0,0 @@
-# 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/>.
-#
-import json
-from typing import Any, Optional
-
-import pydantic
-from pydantic import BaseModel, Field, field_validator
-from typing_extensions import Annotated
-
-from nucliadb_utils.aiopynecone.exceptions import MetadataTooLargeError
-
-KILO_BYTE = 1024
-MAX_METADATA_SIZE = 40 * KILO_BYTE
-MAX_INDEX_NAME_LENGTH = 45
-MAX_VECTOR_ID_LENGTH = 512
-
-
-# Requests
-
-IndexNamePattern = r"^[a-z0-9-]+$"
-
-
-def validate_index_name(value, handler, info):
-    try:
-        return handler(value)
-    except pydantic.ValidationError as e:
-        if any(x["type"] == "string_pattern_mismatch" for x in e.errors()):
-            raise ValueError(
-                f"Invalid field_id: '{value}'. Pinecone index names must be a string with only "
-                "lowercase letters, numbers and dashes."
-            )
-        else:
-            raise e
-
-
-IndexNameStr = Annotated[
-    str,
-    pydantic.StringConstraints(pattern=IndexNamePattern, min_length=1, max_length=MAX_INDEX_NAME_LENGTH),
-    pydantic.WrapValidator(validate_index_name),
-]
-
-
-class CreateIndexRequest(BaseModel):
-    name: IndexNameStr
-    dimension: int
-    metric: str
-    spec: dict[str, Any] = {}
-
-
-class Vector(BaseModel):
-    id: str = Field(min_length=1, max_length=MAX_VECTOR_ID_LENGTH)
-    values: list[float]
-    metadata: dict[str, Any] = {}
-
-    @field_validator("metadata", mode="after")
-    @classmethod
-    def validate_metadata_size(cls, value):
-        json_value = json.dumps(value)
-        if len(json_value) > MAX_METADATA_SIZE:
-            raise MetadataTooLargeError(f"metadata size is too large: {len(json_value)} bytes")
-        return value
-
-
-class UpsertRequest(BaseModel):
-    vectors: list[Vector]
-
-
-# Responses
-
-
-class CreateIndexResponse(BaseModel):
-    host: str
-
-
-class VectorId(BaseModel):
-    id: str
-
-
-class Pagination(BaseModel):
-    next: str
-
-
-class ListResponse(BaseModel):
-    vectors: list[VectorId]
-    pagination: Optional[Pagination] = None
-
-
-class VectorMatch(BaseModel):
-    id: str
-    score: float
-    # Only populated if `includeValues` is set to `True
-    values: Optional[list[float]] = None
-    # Only populated if `includeMetadata` is set to `True
-    metadata: Optional[dict[str, Any]] = None
-
-
-class QueryResponse(BaseModel):
-    matches: list[VectorMatch]
-
-
-class IndexNamespaceStats(BaseModel):
-    vectorCount: int
-
-
-class IndexStats(BaseModel):
-    dimension: int
-    namespaces: dict[str, IndexNamespaceStats] = {}
-    totalVectorCount: int
-
-
-class IndexStatus(BaseModel):
-    ready: bool
-    state: str
-
-
-class IndexDescription(BaseModel):
-    dimension: int
-    host: str
-    metric: str
-    name: str
-    spec: dict[str, Any]
-    status: IndexStatus