rig_lancedb/
lib.rs

1#![cfg_attr(
2    test,
3    allow(
4        clippy::expect_used,
5        clippy::indexing_slicing,
6        clippy::panic,
7        clippy::unwrap_used,
8        clippy::unreachable
9    )
10)]
11//! LanceDB vector store integration for Rig.
12//!
13//! This crate provides [`LanceDbVectorIndex`], a Rig vector store index backed
14//! by LanceDB tables. It supports exact and approximate vector search through
15//! [`SearchType`] and accepts LanceDB SQL filter expressions through
16//! [`LanceDBFilter`].
17//!
18//! The root `rig` facade re-exports this crate as `rig::lancedb` when the
19//! `lancedb` feature is enabled.
20
21use std::ops::Range;
22
23use lancedb::{
24    DistanceType,
25    query::{QueryBase, VectorQuery},
26};
27use rig_core::{
28    embeddings::embedding::EmbeddingModel,
29    vector_store::{
30        VectorStoreError, VectorStoreIndex,
31        request::{FilterError, SearchFilter, VectorSearchRequest},
32    },
33};
34use serde::Deserialize;
35use serde_json::Value;
36use utils::{FilterTableColumns, QueryToJson};
37
38mod utils;
39
40fn lancedb_to_rig_error(e: lancedb::Error) -> VectorStoreError {
41    VectorStoreError::DatastoreError(Box::new(e))
42}
43
44fn serde_to_rig_error(e: serde_json::Error) -> VectorStoreError {
45    VectorStoreError::JsonError(e)
46}
47
48/// Type on which vector searches can be performed for a lanceDb table.
49/// # Example
50/// ```ignore
51/// use rig_lancedb::{LanceDbVectorIndex, SearchParams};
52/// use rig_core::client::ProviderClient;
53/// use rig_core::providers::openai::{Client, TEXT_EMBEDDING_ADA_002, EmbeddingModel};
54///
55/// let openai_client = Client::from_env()?;
56///
57/// let table: lancedb::Table = db.create_table(""); // <-- Replace with your lancedb table here.
58/// let model: EmbeddingModel = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002); // <-- Replace with your embedding model here.
59/// let vector_store_index = LanceDbVectorIndex::new(table, model, "id", SearchParams::default()).await?;
60/// ```
61pub struct LanceDbVectorIndex<M: EmbeddingModel> {
62    /// Defines which model is used to generate embeddings for the vector store.
63    model: M,
64    /// LanceDB table containing embeddings.
65    table: lancedb::Table,
66    /// Column name in `table` that contains the id of a record.
67    id_field: String,
68    /// Vector search params that are used during vector search operations.
69    search_params: SearchParams,
70}
71
72impl<M> LanceDbVectorIndex<M>
73where
74    M: EmbeddingModel,
75{
76    /// Create an instance of `LanceDbVectorIndex` with an existing table and model.
77    /// Define the id field name of the table.
78    /// Define search parameters that will be used to perform vector searches on the table.
79    pub async fn new(
80        table: lancedb::Table,
81        model: M,
82        id_field: &str,
83        search_params: SearchParams,
84    ) -> Result<Self, lancedb::Error> {
85        Ok(Self {
86            table,
87            model,
88            id_field: id_field.to_string(),
89            search_params,
90        })
91    }
92
93    /// Apply the search_params to the vector query.
94    /// This is a helper function used by the methods `top_n` and `top_n_ids` of the `VectorStoreIndex` trait.
95    fn build_query(&self, mut query: VectorQuery) -> VectorQuery {
96        let SearchParams {
97            distance_type,
98            search_type,
99            nprobes,
100            refine_factor,
101            post_filter,
102            column,
103        } = self.search_params.clone();
104
105        if let Some(distance_type) = distance_type {
106            query = query.distance_type(distance_type);
107        }
108
109        if let Some(SearchType::Flat) = search_type {
110            query = query.bypass_vector_index();
111        }
112
113        if let Some(SearchType::Approximate) = search_type {
114            if let Some(nprobes) = nprobes {
115                query = query.nprobes(nprobes);
116            }
117            if let Some(refine_factor) = refine_factor {
118                query = query.refine_factor(refine_factor);
119            }
120        }
121
122        if let Some(true) = post_filter {
123            query = query.postfilter();
124        }
125
126        if let Some(column) = column {
127            query = query.column(column.as_str())
128        }
129
130        query
131    }
132}
133
134/// See [LanceDB vector search](https://lancedb.github.io/lancedb/search/) for more information.
135#[derive(Debug, Clone)]
136pub enum SearchType {
137    // Flat search, also called ENN or kNN.
138    Flat,
139    /// Approximal Nearest Neighbor search, also called ANN.
140    Approximate,
141}
142
143/// An eDSL for filtering expressions, is rendered as a `WHERE` clause
144#[derive(Debug, Clone)]
145pub struct LanceDBFilter(Result<String, FilterError>);
146
147impl serde::Serialize for LanceDBFilter {
148    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
149    where
150        S: serde::Serializer,
151    {
152        match &self.0 {
153            Ok(s) => serializer.serialize_str(s),
154            Err(e) => serializer.collect_str(e),
155        }
156    }
157}
158
159impl<'de> serde::Deserialize<'de> for LanceDBFilter {
160    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
161    where
162        D: serde::Deserializer<'de>,
163    {
164        let s = String::deserialize(deserializer)?;
165        // We can't deserialize to Error, so just create an Ok variant
166        Ok(LanceDBFilter(Ok(s)))
167    }
168}
169
170fn zip_result(
171    l: Result<String, FilterError>,
172    r: Result<String, FilterError>,
173) -> Result<(String, String), FilterError> {
174    l.and_then(|l| r.map(|r| (l, r)))
175}
176
177impl SearchFilter for LanceDBFilter {
178    type Value = serde_json::Value;
179
180    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
181        Self(escape_value(value).map(|s| format!("{} = {s}", key.as_ref())))
182    }
183
184    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
185        Self(escape_value(value).map(|s| format!("{} > {s}", key.as_ref())))
186    }
187
188    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
189        Self(escape_value(value).map(|s| format!("{} < {s}", key.as_ref())))
190    }
191
192    fn and(self, rhs: Self) -> Self {
193        Self(zip_result(self.0, rhs.0).map(|(l, r)| format!("({l}) AND ({r})")))
194    }
195
196    fn or(self, rhs: Self) -> Self {
197        Self(zip_result(self.0, rhs.0).map(|(l, r)| format!("({l}) OR ({r})")))
198    }
199}
200
201fn escape_value(value: serde_json::Value) -> Result<String, FilterError> {
202    use serde_json::Value::*;
203
204    match value {
205        Null => Ok("NULL".into()),
206        Bool(b) => Ok(b.to_string()),
207        Number(n) => Ok(n.to_string()),
208        String(s) => Ok(format!("'{}'", s.replace("'", "''"))),
209        Array(xs) => Ok(format!(
210            "({})",
211            xs.into_iter()
212                .map(escape_value)
213                .collect::<Result<Vec<_>, _>>()?
214                .join(", ")
215        )),
216        Object(_) => Err(FilterError::TypeError(
217            "objects not supported in SQLite backend".into(),
218        )),
219    }
220}
221
222impl LanceDBFilter {
223    pub fn into_inner(self) -> Result<String, FilterError> {
224        self.0
225    }
226
227    #[allow(clippy::should_implement_trait)]
228    pub fn not(self) -> Self {
229        Self(self.0.map(|s| format!("NOT ({s})")))
230    }
231
232    /// IN operator
233    pub fn in_values(key: String, values: Vec<<Self as SearchFilter>::Value>) -> Self {
234        Self(
235            values
236                .into_iter()
237                .map(escape_value)
238                .collect::<Result<Vec<_>, FilterError>>()
239                .map(|xs| xs.join(","))
240                .map(|xs| format!("{key} IN ({xs})")),
241        )
242    }
243
244    /// LIKE operator (string pattern matching)
245    pub fn like<S>(key: String, pattern: S) -> Self
246    where
247        S: AsRef<str>,
248    {
249        Self(
250            escape_value(serde_json::Value::String(pattern.as_ref().into()))
251                .map(|pat| format!("{key} LIKE {pat}")),
252        )
253    }
254
255    /// ILIKE operator (case-insensitive pattern matching)
256    pub fn ilike<S>(key: String, pattern: S) -> Self
257    where
258        S: AsRef<str>,
259    {
260        Self(
261            escape_value(serde_json::Value::String(pattern.as_ref().into()))
262                .map(|pat| format!("{key} ILIKE {pat}")),
263        )
264    }
265
266    /// IS NULL check
267    pub fn is_null(key: String) -> Self {
268        Self(Ok(format!("{key} IS NULL")))
269    }
270
271    /// IS NOT NULL check
272    pub fn is_not_null(key: String) -> Self {
273        Self(Ok(format!("{key} IS NOT NULL")))
274    }
275
276    /// Array has any (for LIST columns with scalar index)
277    pub fn array_has_any(key: String, values: Vec<<Self as SearchFilter>::Value>) -> Self {
278        Self(
279            values
280                .into_iter()
281                .map(escape_value)
282                .collect::<Result<Vec<_>, FilterError>>()
283                .map(|xs| xs.join(","))
284                .map(|xs| format!("array_has_any({key}, ARRAY[{xs}])")),
285        )
286    }
287
288    /// Array has all (for LIST columns with scalar index)
289    pub fn array_has_all(key: String, values: Vec<<Self as SearchFilter>::Value>) -> Self {
290        Self(
291            values
292                .into_iter()
293                .map(escape_value)
294                .collect::<Result<Vec<_>, FilterError>>()
295                .map(|xs| xs.join(","))
296                .map(|xs| format!("array_has_all({key}, ARRAY[{xs}])")),
297        )
298    }
299
300    /// Array length comparison
301    pub fn array_length(key: String, length: i32) -> Self {
302        Self(Ok(format!("array_length({key}) = {length}")))
303    }
304
305    /// BETWEEN operator
306    pub fn between<T>(key: String, Range { start, end }: Range<T>) -> Self
307    where
308        T: PartialOrd + std::fmt::Display + Into<serde_json::Number>,
309    {
310        Self(Ok(format!("{key} BETWEEN {start} AND {end}")))
311    }
312}
313
314/// Parameters used to perform a vector search on a LanceDb table.
315/// # Example
316/// ```
317/// let search_params = rig_lancedb::SearchParams::default().distance_type(lancedb::DistanceType::Cosine);
318/// ```
319#[derive(Debug, Clone, Default)]
320pub struct SearchParams {
321    distance_type: Option<DistanceType>,
322    search_type: Option<SearchType>,
323    nprobes: Option<usize>,
324    refine_factor: Option<u32>,
325    post_filter: Option<bool>,
326    column: Option<String>,
327}
328
329impl SearchParams {
330    /// Sets the distance type of the search params.
331    /// Always set the distance_type to match the value used to train the index.
332    /// The default is DistanceType::L2.
333    pub fn distance_type(mut self, distance_type: DistanceType) -> Self {
334        self.distance_type = Some(distance_type);
335        self
336    }
337
338    /// Sets the search type of the search params.
339    /// By default, ANN will be used if there is an index on the table and kNN will be used if there is NO index on the table.
340    /// To use the mentioned defaults, do not set the search type.
341    pub fn search_type(mut self, search_type: SearchType) -> Self {
342        self.search_type = Some(search_type);
343        self
344    }
345
346    /// Sets the nprobes of the search params.
347    /// Only set this value only when the search type is ANN.
348    /// See [LanceDb ANN Search](https://lancedb.github.io/lancedb/ann_indexes/#querying-an-ann-index) for more information.
349    pub fn nprobes(mut self, nprobes: usize) -> Self {
350        self.nprobes = Some(nprobes);
351        self
352    }
353
354    /// Sets the refine factor of the search params.
355    /// Only set this value only when search type is ANN.
356    /// See [LanceDb ANN Search](https://lancedb.github.io/lancedb/ann_indexes/#querying-an-ann-index) for more information.
357    pub fn refine_factor(mut self, refine_factor: u32) -> Self {
358        self.refine_factor = Some(refine_factor);
359        self
360    }
361
362    /// Sets the post filter of the search params.
363    /// If set to true, filtering will happen after the vector search instead of before.
364    /// See [LanceDb pre/post filtering](https://lancedb.github.io/lancedb/sql/#pre-and-post-filtering) for more information.
365    pub fn post_filter(mut self, post_filter: bool) -> Self {
366        self.post_filter = Some(post_filter);
367        self
368    }
369
370    /// Sets the column of the search params.
371    /// Only set this value if there is more than one column that contains lists of floats.
372    /// If there is only one column of list of floats, this column will be chosen for the vector search automatically.
373    pub fn column(mut self, column: &str) -> Self {
374        self.column = Some(column.to_string());
375        self
376    }
377}
378
379impl<M> VectorStoreIndex for LanceDbVectorIndex<M>
380where
381    M: EmbeddingModel + Sync + Send,
382{
383    type Filter = LanceDBFilter;
384
385    /// Implement the `top_n` method of the `VectorStoreIndex` trait for `LanceDbVectorIndex`.
386    /// # Example
387    /// ```ignore
388    /// use rig_lancedb::{LanceDbVectorIndex, SearchParams};
389    /// use rig_core::client::ProviderClient;
390    /// use rig_core::providers::openai::{EmbeddingModel, Client, TEXT_EMBEDDING_ADA_002};
391    ///
392    /// let openai_client = Client::from_env()?;
393    ///
394    /// let table: lancedb::Table = db.create_table("fake_definitions"); // <-- Replace with your lancedb table here.
395    /// let model: EmbeddingModel = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002); // <-- Replace with your embedding model here.
396    /// let vector_store_index = LanceDbVectorIndex::new(table, model, "id", SearchParams::default()).await?;
397    ///
398    /// // Query the index
399    /// let result = vector_store_index
400    ///     .top_n::<String>("My boss says I zindle too much, what does that mean?", 1)
401    ///     .await?;
402    /// ```
403    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
404        &self,
405        req: VectorSearchRequest<LanceDBFilter>,
406    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
407        let prompt_embedding = self.model.embed_text(req.query()).await?;
408
409        let mut query = self
410            .table
411            .vector_search(prompt_embedding.vec.clone())
412            .map_err(lancedb_to_rig_error)?
413            .limit(req.samples() as usize)
414            .distance_range(None, req.threshold().map(|x| x as f32))
415            .select(lancedb::query::Select::Columns(
416                self.table
417                    .schema()
418                    .await
419                    .map_err(lancedb_to_rig_error)?
420                    .filter_embeddings(),
421            ));
422
423        if let Some(filter) = req.filter() {
424            query = query.only_if(filter.clone().into_inner()?)
425        }
426
427        self.build_query(query)
428            .execute_query()
429            .await?
430            .into_iter()
431            .enumerate()
432            .map(|(i, value)| {
433                Ok((
434                    match value.get("_distance") {
435                        Some(Value::Number(distance)) => distance.as_f64().unwrap_or_default(),
436                        _ => 0.0,
437                    },
438                    match value.get(self.id_field.clone()) {
439                        Some(Value::String(id)) => id.to_string(),
440                        _ => format!("unknown{i}"),
441                    },
442                    serde_json::from_value(value).map_err(serde_to_rig_error)?,
443                ))
444            })
445            .collect()
446    }
447
448    /// Implement the `top_n_ids` method of the `VectorStoreIndex` trait for `LanceDbVectorIndex`.
449    /// # Example
450    /// ```ignore
451    /// use rig_lancedb::{LanceDbVectorIndex, SearchParams};
452    /// use rig_core::client::ProviderClient;
453    /// use rig_core::providers::openai::{Client, TEXT_EMBEDDING_ADA_002, EmbeddingModel};
454    ///
455    /// let openai_client = Client::from_env()?;
456    ///
457    /// let table: lancedb::Table = db.create_table(""); // <-- Replace with your lancedb table here.
458    /// let model: EmbeddingModel = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002); // <-- Replace with your embedding model here.
459    /// let vector_store_index = LanceDbVectorIndex::new(table, model, "id", SearchParams::default()).await?;
460    ///
461    /// // Query the index
462    /// let result = vector_store_index
463    ///     .top_n_ids("My boss says I zindle too much, what does that mean?", 1)
464    ///     .await?;
465    /// ```
466    async fn top_n_ids(
467        &self,
468        req: VectorSearchRequest<LanceDBFilter>,
469    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
470        let prompt_embedding = self.model.embed_text(req.query()).await?;
471
472        let mut query = self
473            .table
474            .query()
475            .select(lancedb::query::Select::Columns(vec![self.id_field.clone()]))
476            .nearest_to(prompt_embedding.vec.clone())
477            .map_err(lancedb_to_rig_error)?
478            .distance_range(None, req.threshold().map(|x| x as f32))
479            .limit(req.samples() as usize);
480
481        if let Some(filter) = req.filter() {
482            query = query.only_if(filter.clone().into_inner()?)
483        }
484
485        self.build_query(query)
486            .execute_query()
487            .await?
488            .into_iter()
489            .map(|value| {
490                Ok((
491                    match value.get("distance") {
492                        Some(Value::Number(distance)) => distance.as_f64().unwrap_or_default(),
493                        _ => 0.0,
494                    },
495                    match value.get(self.id_field.clone()) {
496                        Some(Value::String(id)) => id.to_string(),
497                        _ => "".to_string(),
498                    },
499                ))
500            })
501            .collect()
502    }
503}
rig_lancedb/lib.rs

rig_lancedb/
lib.rs
