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
179                    .id
180                    .clone()
181                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
182                LanceDBDocument {
183                    id: id.clone(),
184                    vector: vec,
185                    content: doc.content,
186                    metadata: doc.metadata,
187                }
188            })
189            .collect();
190
191        let ids: Vec<String> = lancedb_docs.iter().map(|d| d.id.clone()).collect();
192
193        let url = format!("{}/insert", self.table_url());
194        let body = json!({
195            "data": lancedb_docs,
196        });
197
198        let req = self.client.post(&url);
199        let req = self.add_auth(req);
200        let response = req
201            .header("Content-Type", "application/json")
202            .json(&body)
203            .send()
204            .await
205            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
206
207        let status = response.status();
208        if !status.is_success() {
209            let error_text = response.text().await.unwrap_or_default();
210            return Err(VectorStoreError::StorageError(format!(
211                "HTTP {}: {}",
212                status, error_text
213            )));
214        }
215
216        Ok(ids)
217    }
218
219    async fn similarity_search(
220        &self,
221        query_embedding: &[f32],
222        k: usize,
223    ) -> Result<Vec<SearchResult>, VectorStoreError> {
224        let url = format!("{}/search", self.table_url());
225        let body = json!({
226            "vector": query_embedding,
227            "k": k,
228        });
229
230        let req = self.client.post(&url);
231        let req = self.add_auth(req);
232        let response = req
233            .header("Content-Type", "application/json")
234            .json(&body)
235            .send()
236            .await
237            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
238
239        let status = response.status();
240        if !status.is_success() {
241            let error_text = response.text().await.unwrap_or_default();
242            return Err(VectorStoreError::StorageError(format!(
243                "HTTP {}: {}",
244                status, error_text
245            )));
246        }
247
248        let search_response: LanceDBSearchResponse = response
249            .json()
250            .await
251            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
252
253        Ok(search_response
254            .data
255            .into_iter()
256            .map(|item| {
257                let mut doc = Document::new(item.content).with_id(item.id);
258                for (key, value) in item.metadata {
259                    doc = doc.with_metadata(key, value);
260                }
261                SearchResult {
262                    document: doc,
263                    score: item.score.unwrap_or(0.0),
264                }
265            })
266            .collect())
267    }
268
269    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
270        let url = format!("{}/get/{}", self.table_url(), id);
271
272        let req = self.client.get(&url);
273        let req = self.add_auth(req);
274        let response = req
275            .send()
276            .await
277            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
278
279        let status = response.status();
280        if status == reqwest::StatusCode::NOT_FOUND {
281            return Ok(None);
282        }
283        if !status.is_success() {
284            let error_text = response.text().await.unwrap_or_default();
285            return Err(VectorStoreError::StorageError(format!(
286                "HTTP {}: {}",
287                status, error_text
288            )));
289        }
290
291        let item: LanceDBSearchItem = response
292            .json()
293            .await
294            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
295
296        let mut doc = Document::new(item.content).with_id(item.id);
297        for (key, value) in item.metadata {
298            doc = doc.with_metadata(key, value);
299        }
300        Ok(Some(doc))
301    }
302
303    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
304        let url = format!("{}/get/{}", self.table_url(), id);
305
306        let req = self.client.get(&url);
307        let req = self.add_auth(req);
308        let response = req
309            .send()
310            .await
311            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
312
313        let status = response.status();
314        if status == reqwest::StatusCode::NOT_FOUND {
315            return Ok(None);
316        }
317        if !status.is_success() {
318            let error_text = response.text().await.unwrap_or_default();
319            return Err(VectorStoreError::StorageError(format!(
320                "HTTP {}: {}",
321                status, error_text
322            )));
323        }
324
325        let item: LanceDBSearchItem = response
326            .json()
327            .await
328            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
329
330        Ok(Some(item.vector))
331    }
332
333    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
334        let url = format!("{}/delete/{}", self.table_url(), id);
335
336        let req = self.client.delete(&url);
337        let req = self.add_auth(req);
338        let response = req
339            .send()
340            .await
341            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
342
343        let status = response.status();
344        if !status.is_success() {
345            let error_text = response.text().await.unwrap_or_default();
346            return Err(VectorStoreError::StorageError(format!(
347                "HTTP {}: {}",
348                status, error_text
349            )));
350        }
351
352        Ok(())
353    }
354
355    async fn count(&self) -> usize {
356        let url = format!("{}/count", self.table_url());
357
358        let req = self.client.get(&url);
359        let req = self.add_auth(req);
360        let result = req.send().await;
361
362        match result {
363            Ok(response) if response.status().is_success() => {
364                let body: serde_json::Value = response.json().await.unwrap_or_default();
365                body["count"].as_u64().unwrap_or(0) as usize
366            }
367            _ => 0,
368        }
369    }
370
371    async fn clear(&self) -> Result<(), VectorStoreError> {
372        let url = format!("{}/clear", self.table_url());
373
374        let req = self.client.post(&url);
375        let req = self.add_auth(req);
376        let response = req
377            .send()
378            .await
379            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
380
381        let status = response.status();
382        if !status.is_success() {
383            let error_text = response.text().await.unwrap_or_default();
384            return Err(VectorStoreError::StorageError(format!(
385                "HTTP {}: {}",
386                status, error_text
387            )));
388        }
389
390        Ok(())
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    #[test]
399    fn test_config_new() {
400        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
401        assert_eq!(config.uri, "http://localhost:1337");
402        assert_eq!(config.table_name, "my_table");
403        assert!(config.api_key.is_none());
404    }
405
406    #[test]
407    fn test_config_builder() {
408        let config = LanceDBConfig::new("http://localhost:1337", "test")
409            .with_api_key("secret")
410            .with_region("us-east-1");
411        assert_eq!(config.api_key, Some("secret".to_string()));
412        assert_eq!(config.region, Some("us-east-1".to_string()));
413    }
414
415    #[test]
416    fn test_table_url() {
417        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
418        let store = LanceDBVectorStore::new(config);
419        assert_eq!(store.table_url(), "http://localhost:1337/v1/table/my_table");
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!(store.table_url(), "http://localhost:1337/v1/table/my_table");
427    }
428
429    #[test]
430    fn test_store_new() {
431        let config = LanceDBConfig::new("http://localhost:1337", "test");
432        let _store = LanceDBVectorStore::new(config);
433    }
434
435    #[test]
436    fn test_lancedb_document_serialization() {
437        let doc = LanceDBDocument {
438            id: "test-1".to_string(),
439            vector: vec![0.1, 0.2, 0.3],
440            content: "hello world".to_string(),
441            metadata: std::collections::HashMap::new(),
442        };
443        let json = serde_json::to_value(&doc).unwrap();
444        assert_eq!(json["id"], "test-1");
445        assert!(json["vector"].is_array());
446        assert_eq!(json["content"], "hello world");
447    }
448}
lc_vector_stores/lancedb.rs

lc_vector_stores/
lancedb.rs
