lc_vector_stores/
lancedb.rs

1// lc-vector-stores/src/lancedb.rs
2//! LanceDB vector store implementation.
3//!
4//! LanceDB is a serverless, low-latency vector database for AI applications.
5//! This implementation uses the LanceDB HTTP API for remote/server mode.
6//! For embedded/local mode, use the `lancedb` crate directly.
7//!
8//! # Example
9//!
10//! ```rust,ignore
11//! use lc_vector_stores::lancedb::{LanceDBVectorStore, LanceDBConfig};
12//!
13//! let config = LanceDBConfig::new("http://localhost:1337", "my_table");
14//! let store = LanceDBVectorStore::new(config);
15//! store.add_documents(docs, embeddings).await?;
16//! let results = store.similarity_search(&query_embedding, 5).await?;
17//! ```
18
19use async_trait::async_trait;
20use serde::{Deserialize, Serialize};
21use serde_json::json;
22
23use crate::{Document, SearchResult, VectorStore, VectorStoreError};
24
25/// LanceDB configuration.
26#[derive(Debug, Clone)]
27pub struct LanceDBConfig {
28    /// LanceDB server URI (e.g., "http://localhost:1337" or "db://my-db").
29    pub uri: String,
30    /// Table name.
31    pub table_name: String,
32    /// API key (optional, for LanceDB Cloud).
33    pub api_key: Option<String>,
34    /// Region (optional, for LanceDB Cloud).
35    pub region: Option<String>,
36}
37
38impl LanceDBConfig {
39    /// Creates a new LanceDBConfig.
40    pub fn new(uri: impl Into<String>, table_name: impl Into<String>) -> Self {
41        Self {
42            uri: uri.into(),
43            table_name: table_name.into(),
44            api_key: None,
45            region: None,
46        }
47    }
48
49    /// Creates config from environment variables.
50    pub fn from_env_result() -> Result<Self, String> {
51        let uri = std::env::var("LANCEDB_URI")
52            .map_err(|_| "LANCEDB_URI environment variable not set".to_string())?;
53        let table_name = std::env::var("LANCEDB_TABLE_NAME")
54            .map_err(|_| "LANCEDB_TABLE_NAME environment variable not set".to_string())?;
55        let api_key = std::env::var("LANCEDB_API_KEY").ok();
56        let region = std::env::var("LANCEDB_REGION").ok();
57        Ok(Self {
58            uri,
59            table_name,
60            api_key,
61            region,
62        })
63    }
64
65    /// Sets the API key for LanceDB Cloud.
66    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
67        self.api_key = Some(key.into());
68        self
69    }
70
71    /// Sets the region for LanceDB Cloud.
72    pub fn with_region(mut self, region: impl Into<String>) -> Self {
73        self.region = Some(region.into());
74        self
75    }
76}
77
78/// LanceDB vector store.
79///
80/// Uses HTTP API to communicate with LanceDB server.
81pub struct LanceDBVectorStore {
82    config: LanceDBConfig,
83    client: reqwest::Client,
84}
85
86impl std::fmt::Debug for LanceDBVectorStore {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("LanceDBVectorStore")
89            .field("table", &self.config.table_name)
90            .finish()
91    }
92}
93
94impl LanceDBVectorStore {
95    /// Creates a new LanceDBVectorStore with the given configuration.
96    pub fn new(config: LanceDBConfig) -> Self {
97        Self {
98            config,
99            client: reqwest::Client::new(),
100        }
101    }
102
103    /// Creates from environment variables.
104    pub fn from_env_result() -> Result<Self, String> {
105        Ok(Self::new(LanceDBConfig::from_env_result()?))
106    }
107
108    /// Builds the base URL for the table API.
109    fn table_url(&self) -> String {
110        format!(
111            "{}/v1/table/{}",
112            self.config.uri.trim_end_matches('/'),
113            self.config.table_name
114        )
115    }
116
117    /// Adds authorization headers to a request builder.
118    fn add_auth(&self, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
119        let mut req = req;
120        if let Some(ref api_key) = self.config.api_key {
121            req = req.header("x-api-key", api_key);
122        }
123        if let Some(ref region) = self.config.region {
124            req = req.header("x-region", region);
125        }
126        req
127    }
128}
129
130/// Internal document representation for LanceDB.
131#[derive(Debug, Serialize, Deserialize)]
132struct LanceDBDocument {
133    id: String,
134    vector: Vec<f32>,
135    content: String,
136    #[serde(default, skip_serializing_if = "hash_map_is_empty")]
137    metadata: std::collections::HashMap<String, String>,
138}
139
140fn hash_map_is_empty(map: &std::collections::HashMap<String, String>) -> bool {
141    map.is_empty()
142}
143
144/// LanceDB search response.
145#[derive(Debug, Deserialize)]
146struct LanceDBSearchResponse {
147    data: Vec<LanceDBSearchItem>,
148}
149
150#[derive(Debug, Deserialize)]
151struct LanceDBSearchItem {
152    id: String,
153    vector: Vec<f32>,
154    content: String,
155    #[serde(default)]
156    metadata: std::collections::HashMap<String, String>,
157    #[serde(default)]
158    score: Option<f32>,
159}
160
161#[async_trait]
162impl VectorStore for LanceDBVectorStore {
163    async fn add_documents(
164        &self,
165        documents: Vec<Document>,
166        embeddings: Vec<Vec<f32>>,
167    ) -> Result<Vec<String>, VectorStoreError> {
168        if documents.len() != embeddings.len() {
169            return Err(VectorStoreError::EmbeddingError(
170                "Number of documents and embeddings must match".to_string(),
171            ));
172        }
173
174        let lancedb_docs: Vec<LanceDBDocument> = documents
175            .into_iter()
176            .zip(embeddings)
177            .map(|(doc, vec)| {
178                let id = doc.id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
179                LanceDBDocument {
180                    id: id.clone(),
181                    vector: vec,
182                    content: doc.content,
183                    metadata: doc.metadata,
184                }
185            })
186            .collect();
187
188        let ids: Vec<String> = lancedb_docs.iter().map(|d| d.id.clone()).collect();
189
190        let url = format!("{}/insert", self.table_url());
191        let body = json!({
192            "data": lancedb_docs,
193        });
194
195        let req = self.client.post(&url);
196        let req = self.add_auth(req);
197        let response = req
198            .header("Content-Type", "application/json")
199            .json(&body)
200            .send()
201            .await
202            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
203
204        let status = response.status();
205        if !status.is_success() {
206            let error_text = response.text().await.unwrap_or_default();
207            return Err(VectorStoreError::StorageError(format!(
208                "HTTP {}: {}",
209                status, error_text
210            )));
211        }
212
213        Ok(ids)
214    }
215
216    async fn similarity_search(
217        &self,
218        query_embedding: &[f32],
219        k: usize,
220    ) -> Result<Vec<SearchResult>, VectorStoreError> {
221        let url = format!("{}/search", self.table_url());
222        let body = json!({
223            "vector": query_embedding,
224            "k": k,
225        });
226
227        let req = self.client.post(&url);
228        let req = self.add_auth(req);
229        let response = req
230            .header("Content-Type", "application/json")
231            .json(&body)
232            .send()
233            .await
234            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
235
236        let status = response.status();
237        if !status.is_success() {
238            let error_text = response.text().await.unwrap_or_default();
239            return Err(VectorStoreError::StorageError(format!(
240                "HTTP {}: {}",
241                status, error_text
242            )));
243        }
244
245        let search_response: LanceDBSearchResponse = response
246            .json()
247            .await
248            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
249
250        Ok(search_response
251            .data
252            .into_iter()
253            .map(|item| {
254                let mut doc = Document::new(item.content).with_id(item.id);
255                for (key, value) in item.metadata {
256                    doc = doc.with_metadata(key, value);
257                }
258                SearchResult {
259                    document: doc,
260                    score: item.score.unwrap_or(0.0),
261                }
262            })
263            .collect())
264    }
265
266    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
267        let url = format!("{}/get/{}", self.table_url(), id);
268
269        let req = self.client.get(&url);
270        let req = self.add_auth(req);
271        let response = req
272            .send()
273            .await
274            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
275
276        let status = response.status();
277        if status == reqwest::StatusCode::NOT_FOUND {
278            return Ok(None);
279        }
280        if !status.is_success() {
281            let error_text = response.text().await.unwrap_or_default();
282            return Err(VectorStoreError::StorageError(format!(
283                "HTTP {}: {}",
284                status, error_text
285            )));
286        }
287
288        let item: LanceDBSearchItem = response
289            .json()
290            .await
291            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
292
293        let mut doc = Document::new(item.content).with_id(item.id);
294        for (key, value) in item.metadata {
295            doc = doc.with_metadata(key, value);
296        }
297        Ok(Some(doc))
298    }
299
300    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
301        let url = format!("{}/get/{}", self.table_url(), id);
302
303        let req = self.client.get(&url);
304        let req = self.add_auth(req);
305        let response = req
306            .send()
307            .await
308            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
309
310        let status = response.status();
311        if status == reqwest::StatusCode::NOT_FOUND {
312            return Ok(None);
313        }
314        if !status.is_success() {
315            let error_text = response.text().await.unwrap_or_default();
316            return Err(VectorStoreError::StorageError(format!(
317                "HTTP {}: {}",
318                status, error_text
319            )));
320        }
321
322        let item: LanceDBSearchItem = response
323            .json()
324            .await
325            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
326
327        Ok(Some(item.vector))
328    }
329
330    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
331        let url = format!("{}/delete/{}", self.table_url(), id);
332
333        let req = self.client.delete(&url);
334        let req = self.add_auth(req);
335        let response = req
336            .send()
337            .await
338            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
339
340        let status = response.status();
341        if !status.is_success() {
342            let error_text = response.text().await.unwrap_or_default();
343            return Err(VectorStoreError::StorageError(format!(
344                "HTTP {}: {}",
345                status, error_text
346            )));
347        }
348
349        Ok(())
350    }
351
352    async fn count(&self) -> usize {
353        let url = format!("{}/count", self.table_url());
354
355        let req = self.client.get(&url);
356        let req = self.add_auth(req);
357        let result = req.send().await;
358
359        match result {
360            Ok(response) if response.status().is_success() => {
361                let body: serde_json::Value = response.json().await.unwrap_or_default();
362                body["count"].as_u64().unwrap_or(0) as usize
363            }
364            _ => 0,
365        }
366    }
367
368    async fn clear(&self) -> Result<(), VectorStoreError> {
369        let url = format!("{}/clear", self.table_url());
370
371        let req = self.client.post(&url);
372        let req = self.add_auth(req);
373        let response = req
374            .send()
375            .await
376            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
377
378        let status = response.status();
379        if !status.is_success() {
380            let error_text = response.text().await.unwrap_or_default();
381            return Err(VectorStoreError::StorageError(format!(
382                "HTTP {}: {}",
383                status, error_text
384            )));
385        }
386
387        Ok(())
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use super::*;
394
395    #[test]
396    fn test_config_new() {
397        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
398        assert_eq!(config.uri, "http://localhost:1337");
399        assert_eq!(config.table_name, "my_table");
400        assert!(config.api_key.is_none());
401    }
402
403    #[test]
404    fn test_config_builder() {
405        let config = LanceDBConfig::new("http://localhost:1337", "test")
406            .with_api_key("secret")
407            .with_region("us-east-1");
408        assert_eq!(config.api_key, Some("secret".to_string()));
409        assert_eq!(config.region, Some("us-east-1".to_string()));
410    }
411
412    #[test]
413    fn test_table_url() {
414        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
415        let store = LanceDBVectorStore::new(config);
416        assert_eq!(
417            store.table_url(),
418            "http://localhost:1337/v1/table/my_table"
419        );
420    }
421
422    #[test]
423    fn test_table_url_trailing_slash() {
424        let config = LanceDBConfig::new("http://localhost:1337/", "my_table");
425        let store = LanceDBVectorStore::new(config);
426        assert_eq!(
427            store.table_url(),
428            "http://localhost:1337/v1/table/my_table"
429        );
430    }
431
432    #[test]
433    fn test_store_new() {
434        let config = LanceDBConfig::new("http://localhost:1337", "test");
435        let _store = LanceDBVectorStore::new(config);
436    }
437
438    #[test]
439    fn test_lancedb_document_serialization() {
440        let doc = LanceDBDocument {
441            id: "test-1".to_string(),
442            vector: vec![0.1, 0.2, 0.3],
443            content: "hello world".to_string(),
444            metadata: std::collections::HashMap::new(),
445        };
446        let json = serde_json::to_value(&doc).unwrap();
447        assert_eq!(json["id"], "test-1");
448        assert!(json["vector"].is_array());
449        assert_eq!(json["content"], "hello world");
450    }
451}
lc_vector_stores/lancedb.rs

lc_vector_stores/
lancedb.rs
