fix(storage): fix more stuff

Author: o-santi

src/storage/src/storage3/_async/analytics.py

@@ -9,11 +9,11 @@
     SortColumn,
     SortOrder,
 )
-from .request import RequestBuilder
+from .request import AsyncRequestBuilder
 
 
 class AsyncStorageAnalyticsClient:
-    def __init__(self, request: RequestBuilder) -> None:
+    def __init__(self, request: AsyncRequestBuilder) -> None:
         self._request = request
 
     async def create(self, bucket_name: str) -> AnalyticsBucket:

src/storage/src/storage3/_async/request.py

@@ -1,12 +1,14 @@
 from typing import Optional
 
-from httpx import AsyncClient, Headers, QueryParams, Response
+from httpx import AsyncClient, Headers, HTTPStatusError, QueryParams, Response
+from pydantic import ValidationError
 from yarl import URL
 
+from ..exceptions import StorageApiError, VectorBucketErrorMessage
 from ..types import JSON, RequestMethod
 
 
-class RequestBuilder:
+class AsyncRequestBuilder:
     def __init__(self, session: AsyncClient, base_url: URL, headers: Headers) -> None:
         self._session = session
         self._base_url = base_url
@@ -19,12 +21,27 @@ async def send(
         body: JSON = None,
         query_params: Optional[QueryParams] = None,
     ) -> Response:
-        data = await self._session.request(
+        response = await self._session.request(
             method=http_method,
             json=body,
             url=str(self._base_url.joinpath(*path)),
             headers=self.headers,
             params=query_params or QueryParams(),
         )
-        data.raise_for_status()
-        return data
+        try:
+            response.raise_for_status()
+            return response
+        except HTTPStatusError as exc:
+            try:
+                error = VectorBucketErrorMessage.model_validate_json(response.content)
+                raise StorageApiError(
+                    message=error.message,
+                    code=error.code or "400",
+                    status=error.statusCode,
+                ) from exc
+            except ValidationError as exc:
+                raise StorageApiError(
+                    message="The request failed, but could not parse error message response.",
+                    code="LibraryError",
+                    status=response.status_code,
+                ) from exc

src/storage/src/storage3/_async/vectors.py

@@ -5,30 +5,41 @@
 from httpx import AsyncClient, Headers
 from yarl import URL
 
-from ..exceptions import VectorBucketException
+from ..exceptions import StorageApiError, VectorBucketException
 from ..types import (
     JSON,
     DistanceMetric,
+    GetVectorBucketResponse,
+    GetVectorIndexResponse,
     GetVectorsResponse,
-    ListIndexesResponse,
+    ListVectorBucketsResponse,
+    ListVectorIndexesResponse,
     ListVectorsResponse,
     MetadataConfiguration,
     QueryVectorsResponse,
+    VectorBucket,
     VectorData,
     VectorFilter,
     VectorIndex,
+    VectorMatch,
     VectorObject,
 )
-from .request import RequestBuilder
+from .request import AsyncRequestBuilder
+
+
+# used to not send non-required values as `null`
+# for they cannot be null
+def remove_none(**kwargs: JSON) -> JSON:
+    return {key: val for key, val in kwargs.items() if val is not None}
 
 
 class AsyncVectorBucketScope:
-    def __init__(self, request: RequestBuilder, bucket_name: str) -> None:
+    def __init__(self, request: AsyncRequestBuilder, bucket_name: str) -> None:
         self._request = request
         self._bucket_name = bucket_name
 
     def with_metadata(self, **data: JSON) -> JSON:
-        return {"vectorBucketName": self._bucket_name, **data}
+        return remove_none(vectorBucketName=self._bucket_name, **data)
 
     async def create_index(
         self,
@@ -47,26 +58,29 @@ async def create_index(
         )
         await self._request.send(http_method="POST", path=["CreateIndex"], body=body)
 
-    async def get_index(self, index_name: str) -> VectorIndex:
+    async def get_index(self, index_name: str) -> Optional[VectorIndex]:
         body = self.with_metadata(indexName=index_name)
-        data = await self._request.send(
-            http_method="POST", path=["GetIndex"], body=body
-        )
-        return VectorIndex.model_validate(data.content)
+        try:
+            data = await self._request.send(
+                http_method="POST", path=["GetIndex"], body=body
+            )
+            return GetVectorIndexResponse.model_validate_json(data.content).index
+        except StorageApiError:
+            return None
 
     async def list_indexes(
         self,
         next_token: Optional[str] = None,
         max_results: Optional[int] = None,
         prefix: Optional[str] = None,
-    ) -> ListIndexesResponse:
+    ) -> ListVectorIndexesResponse:
         body = self.with_metadata(
             next_token=next_token, max_results=max_results, prefix=prefix
         )
         data = await self._request.send(
             http_method="POST", path=["ListIndexes"], body=body
         )
-        return ListIndexesResponse.model_validate(data.content)
+        return ListVectorIndexesResponse.model_validate_json(data.content)
 
     async def delete_index(self, index_name: str) -> None:
         body = self.with_metadata(indexName=index_name)
@@ -78,33 +92,33 @@ def index(self, index_name: str) -> AsyncVectorIndexScope:
 
 class AsyncVectorIndexScope:
     def __init__(
-        self, request: RequestBuilder, bucket_name: str, index_name: str
+        self, request: AsyncRequestBuilder, bucket_name: str, index_name: str
     ) -> None:
         self._request = request
         self._bucket_name = bucket_name
         self._index_name = index_name
 
     def with_metadata(self, **data: JSON) -> JSON:
-        return {
-            "vectorBucketName": self._bucket_name,
-            "indexName": self._index_name,
+        return remove_none(
+            vectorBucketName=self._bucket_name,
+            indexName=self._index_name,
             **data,
-        }
+        )
 
     async def put(self, vectors: List[VectorObject]) -> None:
         body = self.with_metadata(vectors=[v.as_json() for v in vectors])
         await self._request.send(http_method="POST", path=["PutVectors"], body=body)
 
     async def get(
         self, *keys: str, return_data: bool = True, return_metadata: bool = True
-    ) -> GetVectorsResponse:
+    ) -> List[VectorMatch]:
         body = self.with_metadata(
             keys=keys, returnData=return_data, returnMetadata=return_metadata
         )
         data = await self._request.send(
             http_method="POST", path=["GetVectors"], body=body
         )
-        return GetVectorsResponse.model_validate(data.content)
+        return GetVectorsResponse.model_validate_json(data.content).vectors
 
     async def list(
         self,
@@ -126,7 +140,7 @@ async def list(
         data = await self._request.send(
             http_method="POST", path=["ListVectors"], body=body
         )
-        return ListVectorsResponse.model_validate(data.content)
+        return ListVectorsResponse.model_validate_json(data.content)
 
     async def query(
         self,
@@ -146,7 +160,7 @@ async def query(
         data = await self._request.send(
             http_method="POST", path=["QueryVectors"], body=body
         )
-        return QueryVectorsResponse.model_validate(data.content)
+        return QueryVectorsResponse.model_validate_json(data.content)
 
     async def delete(self, keys: List[str]) -> None:
         if 1 < len(keys) or len(keys) > 500:
@@ -157,7 +171,7 @@ async def delete(self, keys: List[str]) -> None:
 
 class AsyncStorageVectorsClient:
     def __init__(self, url: URL, headers: Headers, session: AsyncClient) -> None:
-        self._request = RequestBuilder(session, base_url=URL(url), headers=headers)
+        self._request = AsyncRequestBuilder(session, base_url=URL(url), headers=headers)
 
     def from_(self, bucket_name: str) -> AsyncVectorBucketScope:
         return AsyncVectorBucketScope(self._request, bucket_name)
@@ -168,7 +182,32 @@ async def create_bucket(self, bucket_name: str) -> None:
             http_method="POST", path=["CreateVectorBucket"], body=body
         )
 
-    # async def get_bucket(self, bucket_name: str) -> GetBucketResponse:
-    #     body = { 'vectorBucketName': bucket_name }
-    #     data = await self._request.send(http_method='POST', path=['GetVectorBucket'], body=body)
-    #     return GetVectorsResponse.model_validate(data.content)
+    async def get_bucket(self, bucket_name: str) -> Optional[VectorBucket]:
+        body = {"vectorBucketName": bucket_name}
+        try:
+            data = await self._request.send(
+                http_method="POST", path=["GetVectorBucket"], body=body
+            )
+            return GetVectorBucketResponse.model_validate_json(
+                data.content
+            ).vectorBucket
+        except StorageApiError:
+            return None
+
+    async def list_buckets(
+        self,
+        prefix: Optional[str] = None,
+        max_results: Optional[int] = None,
+        next_token: Optional[str] = None,
+    ) -> ListVectorBucketsResponse:
+        body = remove_none(prefix=prefix, maxResults=max_results, nextToken=next_token)
+        data = await self._request.send(
+            http_method="POST", path=["ListVectorBuckets"], body=body
+        )
+        return ListVectorBucketsResponse.model_validate_json(data.content)
+
+    async def delete_bucket(self, bucket_name: str) -> None:
+        body = {"vectorBucketName": bucket_name}
+        await self._request.send(
+            http_method="POST", path=["DeleteVectorBucket"], body=body
+        )

src/storage/src/storage3/_sync/vectors.py

@@ -10,7 +10,7 @@
     JSON,
     DistanceMetric,
     GetVectorsResponse,
-    ListIndexesResponse,
+    ListVectorIndexesResponse,
     ListVectorsResponse,
     MetadataConfiguration,
     QueryVectorsResponse,
@@ -55,12 +55,12 @@ def list_indexes(
         next_token: Optional[str] = None,
         max_results: Optional[int] = None,
         prefix: Optional[str] = None,
-    ) -> ListIndexesResponse:
+    ) -> ListVectorIndexesResponse:
         body = self.with_metadata(
             next_token=next_token, max_results=max_results, prefix=prefix
         )
         data = self._request.send(http_method="POST", path=["ListIndexes"], body=body)
-        return ListIndexesResponse.model_validate(data.content)
+        return ListVectorIndexesResponse.model_validate(data.content)
 
     def delete_index(self, index_name: str) -> None:
         body = self.with_metadata(indexName=index_name)

src/storage/src/storage3/exceptions.py

@@ -1,4 +1,6 @@
-from typing import TypedDict
+from typing import Optional, TypedDict, Union
+
+from pydantic import BaseModel
 
 from .utils import StorageException
 
@@ -8,17 +10,24 @@ def __init__(self, msg: str) -> None:
         self.msg = msg
 
 
+class VectorBucketErrorMessage(BaseModel):
+    statusCode: Union[str, int]
+    error: str
+    message: str
+    code: Optional[str] = None
+
+
 class StorageApiErrorDict(TypedDict):
     name: str
     message: str
     code: str
-    status: int
+    status: Union[int, str]
 
 
 class StorageApiError(StorageException):
     """Error raised when an operation on the storage API fails."""
 
-    def __init__(self, message: str, code: str, status: int) -> None:
+    def __init__(self, message: str, code: str, status: Union[int, str]) -> None:
         error_message = (
             f"{{'statusCode': {status}, 'error': {code}, 'message': {message}}}"
         )

src/storage/src/storage3/types.py

@@ -159,13 +159,11 @@ class CreateSignedUploadUrlOptions(BaseModel):
     total=False,
 )
 
-DistanceMetric: TypeAlias = Literal["cosine", "euclidean", "dotproduct"]
+DistanceMetric: TypeAlias = Literal["cosine", "euclidean"]
 
 
 class MetadataConfiguration(BaseModel):
-    non_filterable_metadata_keys: Optional[List[str]] = Field(
-        alias="nonFilterableMetadaKeys", default=None
-    )
+    nonFilterableMetadataKeys: Optional[List[str]]
 
 
 class ListIndexesOptions(BaseModel):
@@ -178,9 +176,9 @@ class ListIndexesResponseItem(BaseModel):
     indexName: str
 
 
-class ListIndexesResponse(BaseModel):
+class ListVectorIndexesResponse(BaseModel):
     indexes: List[ListIndexesResponseItem]
-    nextToken: Optional[str]
+    nextToken: Optional[str] = None
 
 
 class VectorIndex(BaseModel):
@@ -195,6 +193,10 @@ class VectorIndex(BaseModel):
     creation_time: Optional[datetime] = None
 
 
+class GetVectorIndexResponse(BaseModel):
+    index: VectorIndex
+
+
 VectorFilter = Dict[str, Any]
 
 
@@ -205,7 +207,7 @@ class VectorData(BaseModel):
 class VectorObject(BaseModel):
     key: str
     data: VectorData
-    metadata: Optional[dict[str, Any]] = None
+    metadata: Optional[dict[str, Union[str, bool, float]]] = None
 
     def as_json(self) -> JSON:
         return {"key": self.key, "data": dict(self.data), "metadata": self.metadata}
@@ -224,7 +226,7 @@ class GetVectorsResponse(BaseModel):
 
 class ListVectorsResponse(BaseModel):
     vectors: List[VectorMatch]
-    nextToken: Optional[str]
+    nextToken: Optional[str] = None
 
 
 class QueryVectorsResponse(BaseModel):
@@ -247,3 +249,27 @@ class AnalyticsBucket(BaseModel):
 
 class AnalyticsBucketDeleteResponse(BaseModel):
     message: str
+
+
+class VectorBucketEncryptionConfiguration(BaseModel):
+    kmsKeyArn: Optional[str] = None
+    sseType: Optional[str] = None
+
+
+class VectorBucket(BaseModel):
+    vectorBucketName: str
+    creationTime: Optional[datetime] = None
+    encryptionConfiguration: Optional[VectorBucketEncryptionConfiguration] = None
+
+
+class GetVectorBucketResponse(BaseModel):
+    vectorBucket: VectorBucket
+
+
+class ListVectorBucketsItem(BaseModel):
+    vectorBucketName: str
+
+
+class ListVectorBucketsResponse(BaseModel):
+    vectorBuckets: List[ListVectorBucketsItem]
+    nextToken: Optional[str] = None