edgebase_admin/
vectorize.rs

1//! VectorizeClient — Vectorize index access for server-side use.
2//! Note: Vectorize is Edge-only. In local/Docker, the server returns stub responses.
3
4use std::collections::HashMap;
5use std::sync::Arc;
6
7use serde_json::Value;
8
9use crate::generated::admin_api_core::GeneratedAdminApi;
10use edgebase_core::error::Error;
11use edgebase_core::http_client::HttpClient;
12
13/// Client for a user-defined Vectorize index.
14pub struct VectorizeClient {
15    http: Arc<HttpClient>,
16    pub(crate) index: String,
17}
18
19impl VectorizeClient {
20    pub(crate) fn new(http: Arc<HttpClient>, index: &str) -> Self {
21        Self {
22            http,
23            index: index.to_string(),
24        }
25    }
26
27    fn core(&self) -> GeneratedAdminApi<'_> {
28        GeneratedAdminApi::new(&self.http)
29    }
30
31    /// Insert or update vectors.
32    /// Returns mutation result with `ok`, optional `count` and `mutationId`.
33    pub async fn upsert(&self, vectors: &[Value]) -> Result<Value, Error> {
34        let body = serde_json::json!({"action": "upsert", "vectors": vectors});
35        self.core().vectorize_operation(&self.index, &body).await
36    }
37
38    /// Insert vectors (errors on duplicate ID — server returns 409).
39    /// Returns mutation result with `ok`, optional `count` and `mutationId`.
40    pub async fn insert(&self, vectors: &[Value]) -> Result<Value, Error> {
41        let body = serde_json::json!({"action": "insert", "vectors": vectors});
42        self.core().vectorize_operation(&self.index, &body).await
43    }
44
45    /// Search for similar vectors.
46    pub async fn search(
47        &self,
48        vector: &[f64],
49        top_k: usize,
50        filter: Option<&Value>,
51        namespace: Option<&str>,
52        return_values: Option<bool>,
53        return_metadata: Option<&str>,
54    ) -> Result<Vec<HashMap<String, Value>>, Error> {
55        let mut body = serde_json::json!({
56            "action": "search",
57            "vector": vector,
58            "topK": top_k,
59        });
60        if let Some(f) = filter {
61            body["filter"] = f.clone();
62        }
63        if let Some(ns) = namespace {
64            body["namespace"] = serde_json::json!(ns);
65        }
66        if let Some(rv) = return_values {
67            body["returnValues"] = serde_json::json!(rv);
68        }
69        if let Some(rm) = return_metadata {
70            body["returnMetadata"] = serde_json::json!(rm);
71        }
72        let res: Value = self.core().vectorize_operation(&self.index, &body).await?;
73        extract_matches(&res)
74    }
75
76    /// Search by an existing vector's ID (Vectorize v2 only).
77    pub async fn query_by_id(
78        &self,
79        vector_id: &str,
80        top_k: usize,
81        filter: Option<&Value>,
82        namespace: Option<&str>,
83        return_values: Option<bool>,
84        return_metadata: Option<&str>,
85    ) -> Result<Vec<HashMap<String, Value>>, Error> {
86        let mut body = serde_json::json!({
87            "action": "queryById",
88            "vectorId": vector_id,
89            "topK": top_k,
90        });
91        if let Some(f) = filter {
92            body["filter"] = f.clone();
93        }
94        if let Some(ns) = namespace {
95            body["namespace"] = serde_json::json!(ns);
96        }
97        if let Some(rv) = return_values {
98            body["returnValues"] = serde_json::json!(rv);
99        }
100        if let Some(rm) = return_metadata {
101            body["returnMetadata"] = serde_json::json!(rm);
102        }
103        let res: Value = self.core().vectorize_operation(&self.index, &body).await?;
104        extract_matches(&res)
105    }
106
107    /// Retrieve vectors by their IDs.
108    pub async fn get_by_ids(&self, ids: &[&str]) -> Result<Vec<HashMap<String, Value>>, Error> {
109        let body = serde_json::json!({"action": "getByIds", "ids": ids});
110        let res: Value = self.core().vectorize_operation(&self.index, &body).await?;
111        if let Some(vectors) = res.get("vectors").and_then(|v| v.as_array()) {
112            Ok(vectors
113                .iter()
114                .filter_map(|v| {
115                    v.as_object()
116                        .map(|o| o.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
117                })
118                .collect())
119        } else {
120            Ok(vec![])
121        }
122    }
123
124    /// Delete vectors by IDs.
125    /// Returns mutation result with `ok`, optional `count` and `mutationId`.
126    pub async fn delete(&self, ids: &[&str]) -> Result<Value, Error> {
127        let body = serde_json::json!({"action": "delete", "ids": ids});
128        self.core().vectorize_operation(&self.index, &body).await
129    }
130
131    /// Get index info (vector count, dimensions, metric).
132    pub async fn describe(&self) -> Result<Value, Error> {
133        let body = serde_json::json!({"action": "describe"});
134        self.core().vectorize_operation(&self.index, &body).await
135    }
136}
137
138fn extract_matches(res: &Value) -> Result<Vec<HashMap<String, Value>>, Error> {
139    if let Some(matches) = res.get("matches").and_then(|v| v.as_array()) {
140        Ok(matches
141            .iter()
142            .filter_map(|v| {
143                v.as_object()
144                    .map(|o| o.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
145            })
146            .collect())
147    } else {
148        Ok(vec![])
149    }
150}
edgebase_admin/vectorize.rs

edgebase_admin/
vectorize.rs
