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, FilterOp, MetadataFilter, 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    /// POST `/search` 并解析结果(普通与过滤检索共用)。
134    async fn search_impl(
135        &self,
136        body: serde_json::Value,
137    ) -> Result<Vec<SearchResult>, VectorStoreError> {
138        let url = format!("{}/search", self.table_url());
139
140        let req = self.client.post(&url);
141        let req = self.add_auth(req);
142        let response = req
143            .header("Content-Type", "application/json")
144            .json(&body)
145            .send()
146            .await
147            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
148
149        let status = response.status();
150        if !status.is_success() {
151            let error_text = response.text().await.unwrap_or_default();
152            return Err(VectorStoreError::StorageError(format!(
153                "HTTP {}: {}",
154                status, error_text
155            )));
156        }
157
158        let search_response: LanceDBSearchResponse = response
159            .json()
160            .await
161            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
162
163        Ok(search_response
164            .data
165            .into_iter()
166            .map(|item| {
167                let mut doc = Document::new(item.content).with_id(item.id);
168                for (key, value) in item.metadata {
169                    doc = doc.with_metadata(key, value);
170                }
171                SearchResult {
172                    document: doc,
173                    score: item.score.unwrap_or(0.0),
174                }
175            })
176            .collect())
177    }
178}
179
180/// S3: [`MetadataFilter`] → LanceDB SQL `where` 子句。
181///
182/// - 字符串值加单引号并转义内部单引号;数字/布尔原样输出。
183/// - `In`/`Nin` 要求值是数组,生成 `key IN (...)` / `key NOT IN (...)`。
184/// - 无法表达的值类型(对象/嵌套数组用于比较、`In` 的非数组值)返回
185///   [`VectorStoreError::UnsupportedFilter`],不静默忽略。
186pub fn filter_to_sql(filter: &MetadataFilter) -> Result<String, VectorStoreError> {
187    match filter {
188        MetadataFilter::Field { key, op, value } => {
189            let ident = quote_ident(key);
190            match op {
191                FilterOp::In | FilterOp::Nin => {
192                    let arr = value.as_array().ok_or_else(|| {
193                        VectorStoreError::UnsupportedFilter(format!(
194                            "IN/NIN requires an array value for field `{}`, got {}",
195                            key,
196                            value_type(value)
197                        ))
198                    })?;
199                    let items = arr
200                        .iter()
201                        .map(sql_literal)
202                        .collect::<Result<Vec<String>, _>>()?;
203                    let keyword = if matches!(op, FilterOp::In) {
204                        "IN"
205                    } else {
206                        "NOT IN"
207                    };
208                    Ok(format!("{} {} ({})", ident, keyword, items.join(", ")))
209                }
210                _ => Ok(format!("{} {} {}", ident, sql_op(op), sql_literal(value)?)),
211            }
212        }
213        MetadataFilter::And(filters) => join_sql(filters, "AND"),
214        MetadataFilter::Or(filters) => join_sql(filters, "OR"),
215    }
216}
217
218/// AND/OR 组合:每个子过滤加括号后连接。
219fn join_sql(filters: &[MetadataFilter], keyword: &str) -> Result<String, VectorStoreError> {
220    if filters.is_empty() {
221        return Ok("TRUE".to_string());
222    }
223    let parts = filters
224        .iter()
225        .map(filter_to_sql)
226        .collect::<Result<Vec<String>, _>>()?;
227    Ok(parts
228        .iter()
229        .map(|p| format!("({})", p))
230        .collect::<Vec<_>>()
231        .join(&format!(" {} ", keyword)))
232}
233
234/// 标识符(字段名)用双引号包裹并转义内部双引号,防注入。
235fn quote_ident(key: &str) -> String {
236    format!("\"{}\"", key.replace('"', "\"\""))
237}
238
239/// 标量值 → SQL 字面量;无法表达的类型返回 [`VectorStoreError::UnsupportedFilter`]。
240fn sql_literal(value: &serde_json::Value) -> Result<String, VectorStoreError> {
241    match value {
242        serde_json::Value::String(s) => Ok(format!("'{}'", s.replace('\'', "''"))),
243        serde_json::Value::Number(n) => Ok(n.to_string()),
244        serde_json::Value::Bool(b) => Ok(b.to_string()),
245        _ => Err(VectorStoreError::UnsupportedFilter(format!(
246            "cannot translate value of type {} to a SQL literal",
247            value_type(value)
248        ))),
249    }
250}
251
252fn value_type(value: &serde_json::Value) -> &'static str {
253    match value {
254        serde_json::Value::Null => "null",
255        serde_json::Value::Bool(_) => "boolean",
256        serde_json::Value::Number(_) => "number",
257        serde_json::Value::String(_) => "string",
258        serde_json::Value::Array(_) => "array",
259        serde_json::Value::Object(_) => "object",
260    }
261}
262
263fn sql_op(op: &FilterOp) -> &'static str {
264    match op {
265        FilterOp::Eq => "=",
266        FilterOp::Ne => "!=",
267        FilterOp::Gt => ">",
268        FilterOp::Gte => ">=",
269        FilterOp::Lt => "<",
270        FilterOp::Lte => "<=",
271        FilterOp::In | FilterOp::Nin => unreachable!("handled by filter_to_sql"),
272    }
273}
274
275/// Internal document representation for LanceDB.
276#[derive(Debug, Serialize, Deserialize)]
277struct LanceDBDocument {
278    id: String,
279    vector: Vec<f32>,
280    content: String,
281    #[serde(default, skip_serializing_if = "hash_map_is_empty")]
282    metadata: std::collections::HashMap<String, serde_json::Value>,
283}
284
285fn hash_map_is_empty(map: &std::collections::HashMap<String, serde_json::Value>) -> bool {
286    map.is_empty()
287}
288
289/// LanceDB search response.
290#[derive(Debug, Deserialize)]
291struct LanceDBSearchResponse {
292    data: Vec<LanceDBSearchItem>,
293}
294
295#[derive(Debug, Deserialize)]
296struct LanceDBSearchItem {
297    id: String,
298    vector: Vec<f32>,
299    content: String,
300    #[serde(default)]
301    metadata: std::collections::HashMap<String, serde_json::Value>,
302    #[serde(default)]
303    score: Option<f32>,
304}
305
306#[async_trait]
307impl VectorStore for LanceDBVectorStore {
308    async fn add_documents(
309        &self,
310        documents: Vec<Document>,
311        embeddings: Vec<Vec<f32>>,
312    ) -> Result<Vec<String>, VectorStoreError> {
313        if documents.len() != embeddings.len() {
314            return Err(VectorStoreError::EmbeddingError(
315                "Number of documents and embeddings must match".to_string(),
316            ));
317        }
318
319        let lancedb_docs: Vec<LanceDBDocument> = documents
320            .into_iter()
321            .zip(embeddings)
322            .map(|(doc, vec)| {
323                let id = doc
324                    .id
325                    .clone()
326                    .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
327                LanceDBDocument {
328                    id: id.clone(),
329                    vector: vec,
330                    content: doc.content,
331                    metadata: doc.metadata,
332                }
333            })
334            .collect();
335
336        let ids: Vec<String> = lancedb_docs.iter().map(|d| d.id.clone()).collect();
337
338        let url = format!("{}/insert", self.table_url());
339        let body = json!({
340            "data": lancedb_docs,
341        });
342
343        let req = self.client.post(&url);
344        let req = self.add_auth(req);
345        let response = req
346            .header("Content-Type", "application/json")
347            .json(&body)
348            .send()
349            .await
350            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
351
352        let status = response.status();
353        if !status.is_success() {
354            let error_text = response.text().await.unwrap_or_default();
355            return Err(VectorStoreError::StorageError(format!(
356                "HTTP {}: {}",
357                status, error_text
358            )));
359        }
360
361        Ok(ids)
362    }
363
364    async fn similarity_search(
365        &self,
366        query_embedding: &[f32],
367        k: usize,
368    ) -> Result<Vec<SearchResult>, VectorStoreError> {
369        let body = json!({ "vector": query_embedding, "k": k });
370        self.search_impl(body).await
371    }
372
373    /// S3: 带元数据过滤的相似度检索 —— 过滤交给服务端(LanceDB SQL `where` 子句)。
374    async fn similarity_search_with_filter(
375        &self,
376        query_embedding: &[f32],
377        k: usize,
378        filter: Option<&MetadataFilter>,
379    ) -> Result<Vec<SearchResult>, VectorStoreError> {
380        let mut body = json!({ "vector": query_embedding, "k": k });
381        if let Some(f) = filter {
382            // 翻译失败(如 IN 的非数组值/嵌套对象)显式报错,不静默忽略过滤。
383            body["where"] = serde_json::Value::String(filter_to_sql(f)?);
384        }
385        self.search_impl(body).await
386    }
387
388    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
389        let url = format!("{}/get/{}", self.table_url(), id);
390
391        let req = self.client.get(&url);
392        let req = self.add_auth(req);
393        let response = req
394            .send()
395            .await
396            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
397
398        let status = response.status();
399        if status == reqwest::StatusCode::NOT_FOUND {
400            return Ok(None);
401        }
402        if !status.is_success() {
403            let error_text = response.text().await.unwrap_or_default();
404            return Err(VectorStoreError::StorageError(format!(
405                "HTTP {}: {}",
406                status, error_text
407            )));
408        }
409
410        let item: LanceDBSearchItem = response
411            .json()
412            .await
413            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
414
415        let mut doc = Document::new(item.content).with_id(item.id);
416        for (key, value) in item.metadata {
417            doc = doc.with_metadata(key, value);
418        }
419        Ok(Some(doc))
420    }
421
422    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
423        let url = format!("{}/get/{}", self.table_url(), id);
424
425        let req = self.client.get(&url);
426        let req = self.add_auth(req);
427        let response = req
428            .send()
429            .await
430            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
431
432        let status = response.status();
433        if status == reqwest::StatusCode::NOT_FOUND {
434            return Ok(None);
435        }
436        if !status.is_success() {
437            let error_text = response.text().await.unwrap_or_default();
438            return Err(VectorStoreError::StorageError(format!(
439                "HTTP {}: {}",
440                status, error_text
441            )));
442        }
443
444        let item: LanceDBSearchItem = response
445            .json()
446            .await
447            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
448
449        Ok(Some(item.vector))
450    }
451
452    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
453        let url = format!("{}/delete/{}", self.table_url(), id);
454
455        let req = self.client.delete(&url);
456        let req = self.add_auth(req);
457        let response = req
458            .send()
459            .await
460            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
461
462        let status = response.status();
463        if !status.is_success() {
464            let error_text = response.text().await.unwrap_or_default();
465            return Err(VectorStoreError::StorageError(format!(
466                "HTTP {}: {}",
467                status, error_text
468            )));
469        }
470
471        Ok(())
472    }
473
474    async fn count(&self) -> usize {
475        let url = format!("{}/count", self.table_url());
476
477        let req = self.client.get(&url);
478        let req = self.add_auth(req);
479        let result = req.send().await;
480
481        match result {
482            Ok(response) if response.status().is_success() => {
483                let body: serde_json::Value = response.json().await.unwrap_or_default();
484                body["count"].as_u64().unwrap_or(0) as usize
485            }
486            _ => 0,
487        }
488    }
489
490    async fn clear(&self) -> Result<(), VectorStoreError> {
491        let url = format!("{}/clear", self.table_url());
492
493        let req = self.client.post(&url);
494        let req = self.add_auth(req);
495        let response = req
496            .send()
497            .await
498            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
499
500        let status = response.status();
501        if !status.is_success() {
502            let error_text = response.text().await.unwrap_or_default();
503            return Err(VectorStoreError::StorageError(format!(
504                "HTTP {}: {}",
505                status, error_text
506            )));
507        }
508
509        Ok(())
510    }
511}
512
513#[cfg(test)]
514mod tests {
515    use super::*;
516
517    #[test]
518    fn test_config_new() {
519        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
520        assert_eq!(config.uri, "http://localhost:1337");
521        assert_eq!(config.table_name, "my_table");
522        assert!(config.api_key.is_none());
523    }
524
525    #[test]
526    fn test_config_builder() {
527        let config = LanceDBConfig::new("http://localhost:1337", "test")
528            .with_api_key("secret")
529            .with_region("us-east-1");
530        assert_eq!(config.api_key, Some("secret".to_string()));
531        assert_eq!(config.region, Some("us-east-1".to_string()));
532    }
533
534    #[test]
535    fn test_table_url() {
536        let config = LanceDBConfig::new("http://localhost:1337", "my_table");
537        let store = LanceDBVectorStore::new(config);
538        assert_eq!(store.table_url(), "http://localhost:1337/v1/table/my_table");
539    }
540
541    #[test]
542    fn test_table_url_trailing_slash() {
543        let config = LanceDBConfig::new("http://localhost:1337/", "my_table");
544        let store = LanceDBVectorStore::new(config);
545        assert_eq!(store.table_url(), "http://localhost:1337/v1/table/my_table");
546    }
547
548    #[test]
549    fn test_store_new() {
550        let config = LanceDBConfig::new("http://localhost:1337", "test");
551        let _store = LanceDBVectorStore::new(config);
552    }
553
554    #[test]
555    fn test_lancedb_document_serialization() {
556        let doc = LanceDBDocument {
557            id: "test-1".to_string(),
558            vector: vec![0.1, 0.2, 0.3],
559            content: "hello world".to_string(),
560            metadata: std::collections::HashMap::new(),
561        };
562        let json = serde_json::to_value(&doc).unwrap();
563        assert_eq!(json["id"], "test-1");
564        assert!(json["vector"].is_array());
565        assert_eq!(json["content"], "hello world");
566    }
567
568    /// S3: 单字段条件 → SQL 表达式(字符串加引号,数字原样)。
569    #[test]
570    fn test_filter_to_sql_field() {
571        assert_eq!(
572            filter_to_sql(&MetadataFilter::field("lang", FilterOp::Eq, "rust")).unwrap(),
573            r#""lang" = 'rust'"#
574        );
575        assert_eq!(
576            filter_to_sql(&MetadataFilter::field("year", FilterOp::Gte, 2020)).unwrap(),
577            r#""year" >= 2020"#
578        );
579        assert_eq!(
580            filter_to_sql(&MetadataFilter::field("active", FilterOp::Ne, true)).unwrap(),
581            r#""active" != true"#
582        );
583        // 单引号转义
584        assert_eq!(
585            filter_to_sql(&MetadataFilter::field("title", FilterOp::Eq, "it's")).unwrap(),
586            r#""title" = 'it''s'"#
587        );
588    }
589
590    /// S3: IN/NOT IN 需要数组值。
591    #[test]
592    fn test_filter_to_sql_in_nin() {
593        assert_eq!(
594            filter_to_sql(&MetadataFilter::field("tags", FilterOp::In, vec!["a", "b"])).unwrap(),
595            r#""tags" IN ('a', 'b')"#
596        );
597        assert_eq!(
598            filter_to_sql(&MetadataFilter::field("tags", FilterOp::Nin, vec!["x"])).unwrap(),
599            r#""tags" NOT IN ('x')"#
600        );
601        // IN 值非数组 → 显式报错
602        let err = filter_to_sql(&MetadataFilter::field("tags", FilterOp::In, "oops"));
603        assert!(matches!(err, Err(VectorStoreError::UnsupportedFilter(_))));
604    }
605
606    /// S3: AND/OR 组合 → 括号包裹 + 连接词。
607    #[test]
608    fn test_filter_to_sql_and_or() {
609        let f = MetadataFilter::or(vec![
610            MetadataFilter::field("lang", FilterOp::Eq, "python"),
611            MetadataFilter::and(vec![
612                MetadataFilter::field("lang", FilterOp::Eq, "rust"),
613                MetadataFilter::field("year", FilterOp::Gt, 2020),
614            ]),
615        ]);
616        assert_eq!(
617            filter_to_sql(&f).unwrap(),
618            r#"("lang" = 'python') OR (("lang" = 'rust') AND ("year" > 2020))"#
619        );
620    }
621
622    /// S3: 不可表达的标量比较(对象值)显式报错。
623    #[test]
624    fn test_filter_to_sql_unsupported_value() {
625        let f = MetadataFilter::field("nested", FilterOp::Eq, serde_json::json!({ "a": 1 }));
626        assert!(matches!(
627            filter_to_sql(&f),
628            Err(VectorStoreError::UnsupportedFilter(_))
629        ));
630    }
631}
lc_vector_stores/lancedb.rs

lc_vector_stores/
lancedb.rs
