nucliadb-utils 5.0.0.post806__py3-none-any.whl → 5.0.0.post821__py3-none-any.whl

nucliadb_utils/aiopynecone/client.py

@@ -25,10 +25,17 @@ 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 Observer
+from nucliadb_utils.aiopynecone.exceptions import (
+    PineconeAPIError,
+    PineconeRateLimitError,
+    raise_for_status,
+)
 from nucliadb_utils.aiopynecone.models import (
+    CreateIndexRequest,
     CreateIndexResponse,
     ListResponse,
     QueryResponse,
@@ -38,9 +45,13 @@ from nucliadb_utils.aiopynecone.models import (
 
 logger = logging.getLogger(__name__)
 
+
 pinecone_observer = Observer(
     "pinecone_client",
     labels={"type": ""},
+    error_mappings={
+        "rate_limit": PineconeRateLimitError,
+    },
 )
 
 DEFAULT_TIMEOUT = 30
@@ -55,25 +66,18 @@ MAX_UPSERT_PAYLOAD_SIZE = 2 * MEGA_BYTE
 MAX_DELETE_BATCH_SIZE = 1000
 
 
-class PineconeAPIError(Exception):
-    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)
+RETRIABLE_EXCEPTIONS = (
+    PineconeRateLimitError,
+    httpx.ConnectError,
+    httpx.NetworkError,
+)
+
+backoff_handler = backoff.on_exception(
+    backoff.expo,
+    RETRIABLE_EXCEPTIONS,
+    jitter=backoff.random_jitter,
+    max_tries=4,
+)
 
 
 class ControlPlane:
@@ -97,15 +101,17 @@ class ControlPlane:
         Returns:
         - The index host to be used for data plane operations.
         """
-        payload = {
-            "name": name,
-            "dimension": dimension,
-            "metric": metric,
-            "spec": {"serverless": {"cloud": "aws", "region": "us-east-1"}},
-        }
+        payload = CreateIndexRequest(
+            name=name,
+            dimension=dimension,
+            metric=metric,
+            spec={"serverless": {"cloud": "aws", "region": "us-east-1"}},
+        )
         headers = {"Api-Key": self.api_key}
-        http_response = await self.http_session.post("/indexes", json=payload, headers=headers)
-        raise_for_status(http_response)
+        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
 
@@ -121,7 +127,7 @@ class ControlPlane:
         if response.status_code == 404:  # pragma: no cover
             logger.warning("Pinecone index not found.", extra={"index_name": name})
             return
-        raise_for_status(response)
+        raise_for_status("delete_index", response)
 
 
 class DataPlane:
@@ -147,6 +153,7 @@ class DataPlane:
     def _get_request_timeout(self, timeout: Optional[float] = None) -> Optional[float]:
         return timeout or self.client_timeout
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "upsert"})
     async def upsert(self, vectors: list[Vector], timeout: Optional[float] = None) -> None:
         """
@@ -165,7 +172,7 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/upsert", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("upsert", response)
 
     def _estimate_upsert_batch_size(self, vectors: list[Vector]) -> int:
         """
@@ -220,6 +227,7 @@ class DataPlane:
 
         await asyncio.gather(*tasks)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "delete"})
     async def delete(self, ids: list[str], timeout: Optional[float] = None) -> None:
         """
@@ -242,8 +250,9 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/delete", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("delete", response)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "list_page"})
     async def list_page(
         self,
@@ -279,7 +288,7 @@ class DataPlane:
             "/vectors/list",
             **post_kwargs,
         )
-        raise_for_status(response)
+        raise_for_status("list_page", response)
         return ListResponse.model_validate(response.json())
 
     async def list_all(
@@ -306,6 +315,7 @@ class DataPlane:
                 break
             pagination_token = response.pagination.next
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "delete_all"})
     async def delete_all(self, timeout: Optional[float] = None):
         """
@@ -324,7 +334,7 @@ class DataPlane:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/vectors/delete", **post_kwargs)
         try:
-            raise_for_status(response)
+            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.
@@ -366,6 +376,7 @@ class DataPlane:
 
         await asyncio.gather(*tasks)
 
+    @backoff_handler
     @pinecone_observer.wrap({"type": "query"})
     async def query(
         self,
@@ -404,7 +415,7 @@ class DataPlane:
         if request_timeout is not None:
             post_kwargs["timeout"] = timeout
         response = await self.http_session.post("/query", **post_kwargs)
-        raise_for_status(response)
+        raise_for_status("query", response)
         return QueryResponse.model_validate(response.json())
 
 
@@ -459,28 +470,6 @@ class PineconeSession:
         return DataPlane(api_key=api_key, index_host_session=index_host_session, timeout=timeout)
 
 
-def raise_for_status(response: httpx.Response):
-    try:
-        response.raise_for_status()
-    except httpx.HTTPStatusError:
-        code = None
-        message = None
-        details = None
-        try:
-            resp_json = response.json()
-            code = resp_json.get("code")
-            message = resp_json.get("message")
-            details = resp_json.get("details")
-        except Exception:
-            message = response.text
-        raise PineconeAPIError(
-            http_status_code=response.status_code,
-            code=code,
-            message=message,
-            details=details,
-        )
-
-
 def batchify(iterable: Iterable, batch_size: int):
     """
     Split an iterable into batches of batch_size

nucliadb_utils/aiopynecone/exceptions.py

@@ -0,0 +1,91 @@
+# 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": ""})
+
+
+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 PineconeRateLimitError(PineconeAPIError):
+    """
+    Raised when the client has exceeded the rate limit to be able to backoff and retry.
+    """
+
+    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})
+        code = None
+        message = None
+        details = None
+        try:
+            resp_json = response.json()
+            error = resp_json.get("error") or {}
+            code = error.get("code")
+            message = error.get("message")
+            details = error.get("details")
+        except Exception:  # pragma: no cover
+            message = response.text
+        if response.status_code == 429:
+            raise PineconeRateLimitError(
+                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

@@ -20,14 +20,47 @@
 import json
 from typing import Any, Optional
 
+import pydantic
 from pydantic import BaseModel, Field, field_validator
+from typing_extensions import Annotated
 
 KILO_BYTE = 1024
 MAX_METADATA_SIZE = 40 * KILO_BYTE
+MAX_INDEX_NAME_LENGTH = 45
 
 
 # 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),
+    pydantic.StringConstraints(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(max_length=512)

{nucliadb_utils-5.0.0.post806.dist-info → nucliadb_utils-5.0.0.post821.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 5.0.0.post806
+Version: 5.0.0.post821
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -23,8 +23,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt >=2.4.0
 Requires-Dist: memorylru >=1.1.2
 Requires-Dist: mrflagly
-Requires-Dist: nucliadb-protos >=5.0.0.post806
-Requires-Dist: nucliadb-telemetry >=5.0.0.post806
+Requires-Dist: nucliadb-protos >=5.0.0.post821
+Requires-Dist: nucliadb-telemetry >=5.0.0.post821
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-5.0.0.post806.dist-info → nucliadb_utils-5.0.0.post821.dist-info}/RECORD

@@ -18,8 +18,9 @@ nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=mwcI3aIHAvU5KOGqd_Uz_d1XQzXhk_-NWY8NqU1lfb0,7307
 nucliadb_utils/utilities.py,sha256=oz3tEODG2g3todnyvA-nW1Ou6xXDveL_tMKTDGdWXM4,15287
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb_utils/aiopynecone/client.py,sha256=kvGLCzSbDaMIWx0LK9WBAL7QsyuPmUjC7cPv5djMdFw,19028
-nucliadb_utils/aiopynecone/models.py,sha256=DVlCVrinHAaDxuii3fzdciFn4dqS5HTTyBavnoDlR2U,2024
+nucliadb_utils/aiopynecone/client.py,sha256=wUQIUZHKvhMhmLVfwrOF_nMBcf9l-4mXCvlSI0l0H24,18472
+nucliadb_utils/aiopynecone/exceptions.py,sha256=hFhq-UEY4slqNWjObXr_LPnRf_AQ1vpcG4SF2XRFd1E,2873
+nucliadb_utils/aiopynecone/models.py,sha256=sEmifzQ6rvqIB8nbkJbh8-hrCW4j8J9lJ_xLRaTAqro,2934
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=dn5ZnCVQUlCcvdjzaORghbrjk9QgVGrtkfIftq30Bp8,2819
 nucliadb_utils/audit/basic.py,sha256=NViey6mKbCXqRTLDBX2xNTcCg9I-2e4oB2xkekuhDvM,3392
@@ -63,8 +64,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=c3gZJJWmvOftruJkIQIwB3q_hh3uxEhqGIAVWim1Bbk,1343
 nucliadb_utils/tests/nats.py,sha256=Tosonm9A9cusImyji80G4pgdXEHNVPaCLT5TbFK_ra0,7543
 nucliadb_utils/tests/s3.py,sha256=YB8QqDaBXxyhHonEHmeBbRRDmvB7sTOaKBSi8KBGokg,2330
-nucliadb_utils-5.0.0.post806.dist-info/METADATA,sha256=HzorldEg8W_vDYIX6r7rQgVf_MaXkNeJQQNdrYgSblc,2073
-nucliadb_utils-5.0.0.post806.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
-nucliadb_utils-5.0.0.post806.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.0.0.post806.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.0.0.post806.dist-info/RECORD,,
+nucliadb_utils-5.0.0.post821.dist-info/METADATA,sha256=54y8hOk8DjtRFOBXlvQnqvcvOOgLvdV3LKexpyS-iMM,2073
+nucliadb_utils-5.0.0.post821.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
+nucliadb_utils-5.0.0.post821.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.0.0.post821.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.0.0.post821.dist-info/RECORD,,