From ca41a868154eeb4b3b6cc1a88344a6099e70d814 Mon Sep 17 00:00:00 2001
From: Garvit Gupta <garvit@cloudflare.com>
Date: Tue, 22 Oct 2024 19:09:23 +0530
Subject: [PATCH] Refactor according to Vectorize V2 spec

---
 worker/src/vectorize.rs | 44 +++++++++++++++++++----------------------
 1 file changed, 20 insertions(+), 24 deletions(-)

diff --git a/worker/src/vectorize.rs b/worker/src/vectorize.rs
index ad954d68..a49c2e51 100644
--- a/worker/src/vectorize.rs
+++ b/worker/src/vectorize.rs
@@ -32,27 +32,18 @@ pub enum VectorizeIndexConfig {
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 /// Metadata about an existing index.
-///
-/// This type is exclusively for the Vectorize **beta** and will be deprecated once Vectorize RC is released.
 pub struct VectorizeIndexDetails {
-    pub id: String,
     pub name: String,
     pub description: Option<String>,
     pub config: VectorizeIndexConfig,
-    pub vectors_count: u64,
 }
 
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 /// Results of an operation that performed a mutation on a set of vectors.
-/// Here, `ids` is a list of vectors that were successfully processed.
-///
-/// This type is exclusively for the Vectorize **beta** and will be deprecated once Vectorize RC is released.
-pub struct VectorizeVectorMutation {
-    /// List of ids of vectors that were successfully processed.
-    pub ids: Vec<String>,
-    /// Total count of the number of processed vectors.
-    pub count: u64,
+pub struct VectorizeVectorAsyncMutation {
+    /// The unique identifier for the async mutation operation containing the changeset.
+    pub mutation_id: String,
 }
 
 #[derive(Debug, Serialize)]
@@ -90,7 +81,7 @@ impl<'a> VectorizeVector<'a> {
     }
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Debug, Default, Serialize)]
 #[serde(rename_all = "kebab-case")]
 /// Metadata return levels for a Vectorize query.
 pub enum VectorizeMetadataRetrievalLevel {
@@ -99,6 +90,7 @@ pub enum VectorizeMetadataRetrievalLevel {
     /// Return all metadata fields configured for indexing in the vector return set. This level of retrieval is "free" in that no additional overhead is incurred returning this data. However, note that indexed metadata is subject to truncation (especially for larger strings).
     Indexed,
     /// No indexed metadata will be returned.
+    #[default]
     None,
 }
 
@@ -120,13 +112,14 @@ type VectorizeVectorMetadataFilter =
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct VectorizeQueryOptions {
-    // Default 3, max 20
+    // Default 5, max 100
     top_k: u8,
+    /// Return vectors from the specified namespace. Default `none`.
     namespace: Option<String>,
     /// Return vector values. Default `false`.
     return_values: bool,
-    /// Return vector metadata. Default `false`.
-    return_metadata: bool,
+    /// Return vector metadata. Default `None`.
+    return_metadata: VectorizeMetadataRetrievalLevel,
     /// Default `none`.
     filter: Option<VectorizeVectorMetadataFilter>,
 }
@@ -151,7 +144,10 @@ impl VectorizeQueryOptions {
         self
     }
 
-    pub fn with_return_metadata(mut self, return_metadata: bool) -> Self {
+    pub fn with_return_metadata(
+        mut self,
+        return_metadata: VectorizeMetadataRetrievalLevel,
+    ) -> Self {
         self.return_metadata = return_metadata;
         self
     }
@@ -173,10 +169,10 @@ impl VectorizeQueryOptions {
 impl Default for VectorizeQueryOptions {
     fn default() -> Self {
         Self {
-            top_k: 3,
+            top_k: 5,
             namespace: None,
             return_values: false,
-            return_metadata: false,
+            return_metadata: VectorizeMetadataRetrievalLevel::None,
             filter: None,
         }
     }
@@ -191,6 +187,8 @@ pub struct VectorizeVectorResult {
     pub values: Option<Vec<f32>>,
     /// Metadata associated with the vector. Includes the values of other fields and potentially additional details.
     pub metadata: Option<serde_json::Map<String, serde_json::Value>>,
+    /** The namespace the vector belongs to. */
+    pub namespace: Option<String>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -209,8 +207,6 @@ pub struct VectorizeMatches {
 }
 
 /// A Vectorize Vector Search Index for querying vectors/embeddings.
-///
-/// This type is exclusively for the Vectorize **beta** and will be deprecated once Vectorize RC is released.
 pub struct VectorizeIndex(VectorizeIndexSys);
 
 unsafe impl Send for VectorizeIndex {}
@@ -233,7 +229,7 @@ impl VectorizeIndex {
     pub async fn insert<'a>(
         &self,
         vectors: &[VectorizeVector<'a>],
-    ) -> Result<VectorizeVectorMutation> {
+    ) -> Result<VectorizeVectorAsyncMutation> {
         let promise = self
             .0
             .insert(serde_wasm_bindgen::to_value(&vectors)?.into())?;
@@ -246,7 +242,7 @@ impl VectorizeIndex {
     pub async fn upsert<'a>(
         &self,
         vectors: &[VectorizeVector<'a>],
-    ) -> Result<VectorizeVectorMutation> {
+    ) -> Result<VectorizeVectorAsyncMutation> {
         let promise = self
             .0
             .upsert(serde_wasm_bindgen::to_value(&vectors)?.into())?;
@@ -269,7 +265,7 @@ impl VectorizeIndex {
     }
 
     /// Delete a list of vectors with a matching id.
-    pub async fn delete_by_ids<'a, T>(&self, ids: T) -> Result<VectorizeVectorMutation>
+    pub async fn delete_by_ids<'a, T>(&self, ids: T) -> Result<VectorizeVectorAsyncMutation>
     where
         T: IntoIterator<Item = &'a str>,
     {