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

lc_vector_stores/
lancedb.rs
