fix(storage): misc fixes to serialization of basemodels

Author: o-santi

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

@@ -54,7 +54,9 @@ async def create_index(
             dimension=dimension,
             distanceMetric=distance_metric,
             dataType=data_type,
-            metadataConfiguration=dict(metadata) if metadata else None,
+            metadataConfiguration=metadata.model_dump(by_alias=True)
+            if metadata
+            else None,
         )
         await self._request.send(http_method="POST", path=["CreateIndex"], body=body)
 
@@ -106,7 +108,9 @@ def with_metadata(self, **data: JSON) -> JSON:
         )
 
     async def put(self, vectors: List[VectorObject]) -> None:
-        body = self.with_metadata(vectors=[v.as_json() for v in vectors])
+        body = self.with_metadata(
+            vectors=[v.model_dump(exclude_none=True) for v in vectors]
+        )
         await self._request.send(http_method="POST", path=["PutVectors"], body=body)
 
     async def get(
@@ -149,7 +153,7 @@ async def query(
         filter: Optional[VectorFilter] = None,
         return_distance: bool = True,
         return_metadata: bool = True,
-    ) -> QueryVectorsResponse:
+    ) -> List[VectorMatch]:
         body = self.with_metadata(
             queryVector=dict(query_vector),
             topK=topK,
@@ -160,10 +164,10 @@ async def query(
         data = await self._request.send(
             http_method="POST", path=["QueryVectors"], body=body
         )
-        return QueryVectorsResponse.model_validate_json(data.content)
+        return QueryVectorsResponse.model_validate_json(data.content).vectors
 
     async def delete(self, keys: List[str]) -> None:
-        if 1 < len(keys) or len(keys) > 500:
+        if len(keys) < 1 or len(keys) > 500:
             raise VectorBucketException("Keys batch size must be between 1 and 500.")
         body = self.with_metadata(keys=keys)
         await self._request.send(http_method="POST", path=["DeleteVectors"], body=body)

src/storage/src/storage3/types.py

@@ -207,14 +207,11 @@ class VectorObject(BaseModel, extra=Extra.ignore):
     data: VectorData
     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}
-
 
 class VectorMatch(BaseModel, extra=Extra.ignore):
     key: str
     data: Optional[VectorData] = None
-    distance: Optional[int] = None
+    distance: Optional[float] = None
     metadata: Optional[dict[str, Any]] = None
 
 
@@ -228,7 +225,7 @@ class ListVectorsResponse(BaseModel, extra=Extra.ignore):
 
 
 class QueryVectorsResponse(BaseModel, extra=Extra.ignore):
-    matches: List[VectorMatch]
+    vectors: List[VectorMatch]
 
 
 class AnalyticsBucket(BaseModel, extra=Extra.ignore):