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

rig_lancedb/
lib.rs
