rig_mongodb/
lib.rs

1//! MongoDB vector store integration for Rig.
2//!
3//! This crate provides [`MongoDbVectorIndex`], a Rig vector store index backed
4//! by MongoDB Atlas Vector Search or compatible MongoDB vector search indexes.
5//!
6//! The root `rig` facade re-exports this crate as `rig::mongodb` when the
7//! `mongodb` feature is enabled.
8
9use futures::StreamExt;
10use mongodb::bson::{self, Bson, Document, doc, to_bson};
11
12use rig_core::{
13    Embed,
14    embeddings::embedding::{Embedding, EmbeddingModel},
15    vector_store::{
16        InsertDocuments, VectorStoreError, VectorStoreIndex,
17        request::{DynamicSearchFilter, Filter, FilterError, SearchFilter, VectorSearchRequest},
18    },
19};
20use serde::{Deserialize, Serialize};
21
22#[derive(Debug, Serialize, Deserialize)]
23#[serde(rename_all = "camelCase")]
24struct SearchIndex {
25    id: String,
26    name: String,
27    #[serde(rename = "type")]
28    index_type: String,
29    status: String,
30    queryable: bool,
31    latest_definition: LatestDefinition,
32}
33
34impl SearchIndex {
35    async fn get_search_index<C: Send + Sync>(
36        collection: mongodb::Collection<C>,
37        index_name: &str,
38    ) -> Result<SearchIndex, VectorStoreError> {
39        collection
40            .list_search_indexes()
41            .name(index_name)
42            .await
43            .map_err(VectorStoreError::datastore)?
44            .with_type::<SearchIndex>()
45            .next()
46            .await
47            .transpose()
48            .map_err(VectorStoreError::datastore)?
49            .ok_or(VectorStoreError::DatastoreError("Index not found".into()))
50    }
51}
52
53#[derive(Debug, Serialize, Deserialize)]
54struct LatestDefinition {
55    fields: Vec<Field>,
56}
57
58#[derive(Debug, Serialize, Deserialize)]
59#[serde(rename_all = "camelCase")]
60struct Field {
61    #[serde(rename = "type")]
62    field_type: String,
63    path: String,
64    num_dimensions: i32,
65    similarity: String,
66}
67
68/// A vector index for a MongoDB collection.
69/// # Example
70/// ```no_run
71/// use rig_mongodb::{MongoDbVectorIndex, SearchParams};
72/// use rig_core::{providers::openai, vector_store::{VectorStoreIndex, VectorSearchRequest}, client::{ProviderClient, EmbeddingsClient}};
73///
74/// # async fn example() -> anyhow::Result<()> {
75/// #[derive(serde::Deserialize, serde::Serialize, Debug)]
76/// struct WordDefinition {
77///     #[serde(rename = "_id")]
78///     id: String,
79///     definition: String,
80///     embedding: Vec<f64>,
81/// }
82///
83/// let mongodb_client = mongodb::Client::with_uri_str("mongodb://localhost:27017").await?; // <-- replace with your mongodb uri.
84/// let openai_client = openai::Client::from_env()?;
85///
86/// let collection = mongodb_client.database("db").collection::<WordDefinition>(""); // <-- replace with your mongodb collection.
87///
88/// let model = openai_client.embedding_model(openai::TEXT_EMBEDDING_ADA_002); // <-- replace with your embedding model.
89/// let index = MongoDbVectorIndex::new(
90///     collection,
91///     model,
92///     "vector_index", // <-- replace with the name of the index in your mongodb collection.
93///     SearchParams::new(), // <-- field name in `Document` that contains the embeddings.
94/// )
95/// .await?;
96///
97/// let req = VectorSearchRequest::builder()
98///     .query("My boss says I zindle too much, what does that mean?")
99///     .samples(1)
100///     .build();
101///
102/// // Query the index
103/// let definitions = index
104///     .top_n::<WordDefinition>(req)
105///     .await?;
106/// # Ok(())
107/// # }
108/// # let _ = example();
109/// ```
110pub struct MongoDbVectorIndex<C, M>
111where
112    C: Send + Sync,
113    M: EmbeddingModel,
114{
115    collection: mongodb::Collection<C>,
116    model: M,
117    index_name: String,
118    embedded_field: String,
119    search_params: SearchParams,
120}
121
122impl<C, M> MongoDbVectorIndex<C, M>
123where
124    C: Send + Sync,
125    M: EmbeddingModel,
126{
127    /// Vector search stage of aggregation pipeline of mongoDB collection.
128    /// To be used by implementations of top_n and top_n_ids methods on VectorStoreIndex trait for MongoDbVectorIndex.
129    fn pipeline_search_stage(
130        &self,
131        prompt_embedding: &Embedding,
132        req: &VectorSearchRequest<MongoDbSearchFilter>,
133    ) -> bson::Document {
134        let SearchParams {
135            exact,
136            num_candidates,
137        } = &self.search_params;
138
139        let samples = req.samples() as usize;
140
141        let thresh = req
142            .threshold()
143            .map(|thresh| MongoDbSearchFilter::gte("score".into(), thresh.into()));
144
145        let filter = match (thresh, req.filter()) {
146            (Some(thresh), Some(filt)) => thresh.and(filt.clone()).into_inner(),
147            (Some(thresh), _) => thresh.into_inner(),
148            (_, Some(filt)) => filt.clone().into_inner(),
149            _ => Default::default(),
150        };
151
152        doc! {
153          "$vectorSearch": {
154            "index": &self.index_name,
155            "path": self.embedded_field.clone(),
156            "queryVector": &prompt_embedding.vec,
157            "numCandidates": num_candidates.unwrap_or((samples * 10) as u32),
158            "limit": samples as u32,
159            "filter": filter,
160            "exact": exact.unwrap_or(false)
161          }
162        }
163    }
164
165    /// Embeds the query, runs the vector-search aggregation pipeline with the
166    /// given `$project` stage, and extracts `(score, id, document)` per row.
167    async fn run_search_pipeline(
168        &self,
169        req: &VectorSearchRequest<MongoDbSearchFilter>,
170        project_stage: bson::Document,
171    ) -> Result<Vec<(f64, String, serde_json::Value)>, VectorStoreError> {
172        let prompt_embedding = self.model.embed_text(req.query()).await?;
173
174        let pipeline = vec![
175            self.pipeline_search_stage(&prompt_embedding, req),
176            self.pipeline_score_stage(),
177            project_stage,
178        ];
179
180        let mut cursor = self
181            .collection
182            .aggregate(pipeline)
183            .await
184            .map_err(VectorStoreError::datastore)?
185            .with_type::<serde_json::Value>();
186
187        let mut results = Vec::new();
188        while let Some(doc) = cursor.next().await {
189            let doc = doc.map_err(VectorStoreError::datastore)?;
190            let score = doc
191                .get("score")
192                .and_then(serde_json::Value::as_f64)
193                .ok_or_else(|| {
194                    VectorStoreError::DatastoreError(
195                        "MongoDB vector search result missing numeric score".into(),
196                    )
197                })?;
198            let id = doc
199                .get("_id")
200                .ok_or_else(|| {
201                    VectorStoreError::DatastoreError(
202                        "MongoDB vector search result missing _id".into(),
203                    )
204                })?
205                .to_string();
206            results.push((score, id, doc));
207        }
208
209        tracing::info!(target: "rig",
210            "Selected documents: {}",
211            results.iter()
212                .map(|(distance, id, _)| format!("{id} ({distance})"))
213                .collect::<Vec<String>>()
214                .join(", ")
215        );
216
217        Ok(results)
218    }
219
220    /// Score declaration stage of aggregation pipeline of mongoDB collection.
221    /// /// To be used by implementations of top_n and top_n_ids methods on VectorStoreIndex trait for MongoDbVectorIndex.
222    fn pipeline_score_stage(&self) -> bson::Document {
223        doc! {
224          "$addFields": {
225            "score": { "$meta": "vectorSearchScore" }
226          }
227        }
228    }
229}
230
231impl<C, M> MongoDbVectorIndex<C, M>
232where
233    M: EmbeddingModel,
234    C: Send + Sync,
235{
236    /// Create a new `MongoDbVectorIndex`.
237    ///
238    /// The index (of type "vector") must already exist for the MongoDB collection.
239    /// See the MongoDB [documentation](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/) for more information on creating indexes.
240    pub async fn new(
241        collection: mongodb::Collection<C>,
242        model: M,
243        index_name: &str,
244        search_params: SearchParams,
245    ) -> Result<Self, VectorStoreError> {
246        let search_index = SearchIndex::get_search_index(collection.clone(), index_name).await?;
247
248        if !search_index.queryable {
249            return Err(VectorStoreError::DatastoreError(
250                "Index is not queryable".into(),
251            ));
252        }
253
254        let embedded_field = search_index
255            .latest_definition
256            .fields
257            .into_iter()
258            .map(|field| field.path)
259            .next()
260            // This error shouldn't occur if the index is queryable
261            .ok_or(VectorStoreError::DatastoreError(
262                "No embedded fields found".into(),
263            ))?;
264
265        Ok(Self {
266            collection,
267            model,
268            index_name: index_name.to_string(),
269            embedded_field,
270            search_params,
271        })
272    }
273}
274
275/// See [MongoDB Vector Search](`https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/`) for more information
276/// on each of the fields
277#[derive(Default)]
278pub struct SearchParams {
279    exact: Option<bool>,
280    num_candidates: Option<u32>,
281}
282
283impl SearchParams {
284    /// Initializes a new `SearchParams` with default values.
285    pub fn new() -> Self {
286        Self {
287            exact: None,
288            num_candidates: None,
289        }
290    }
291
292    /// Sets the exact field of the search params.
293    /// If exact is true, an ENN vector search will be performed, otherwise, an ANN search will be performed.
294    /// By default, exact is false.
295    /// See [MongoDB vector Search](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/) for more information.
296    pub fn exact(mut self, exact: bool) -> Self {
297        self.exact = Some(exact);
298        self
299    }
300
301    /// Sets the num_candidates field of the search params.
302    /// Only set this field if exact is set to false.
303    /// Number of nearest neighbors to use during the search.
304    /// See [MongoDB vector Search](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/) for more information.
305    pub fn num_candidates(mut self, num_candidates: u32) -> Self {
306        self.num_candidates = Some(num_candidates);
307        self
308    }
309}
310
311#[derive(Clone, Debug, Serialize, Deserialize)]
312pub struct MongoDbSearchFilter(Document);
313
314impl SearchFilter for MongoDbSearchFilter {
315    type Value = Bson;
316
317    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
318        let key = key.as_ref().to_owned();
319        Self(doc! { key: value })
320    }
321
322    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
323        let key = key.as_ref().to_owned();
324        Self(doc! { key: { "$gt": value } })
325    }
326
327    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
328        let key = key.as_ref().to_owned();
329        Self(doc! { key: { "$lt": value } })
330    }
331
332    fn and(self, rhs: Self) -> Self {
333        Self(doc! { "$and": [ self.0, rhs.0 ]})
334    }
335
336    fn or(self, rhs: Self) -> Self {
337        Self(doc! { "$or": [ self.0, rhs.0 ]})
338    }
339}
340
341impl MongoDbSearchFilter {
342    fn into_inner(self) -> Document {
343        self.0
344    }
345
346    pub fn gte(key: String, value: <Self as SearchFilter>::Value) -> Self {
347        Self(doc! { key: { "$gte": value } })
348    }
349
350    pub fn lte(key: String, value: <Self as SearchFilter>::Value) -> Self {
351        Self(doc! { key: { "$lte": value } })
352    }
353
354    #[allow(clippy::should_implement_trait)]
355    pub fn not(self) -> Self {
356        Self(doc! { "$nor": [self.0] })
357    }
358
359    /// Tests whether the value at `key` is the BSON type `typ`
360    pub fn is_type(key: String, typ: &'static str) -> Self {
361        Self(doc! { key: { "$type": typ } })
362    }
363
364    pub fn size(key: String, size: i32) -> Self {
365        Self(doc! { key: { "$size": size } })
366    }
367
368    // Array ops
369    pub fn all(key: String, values: Vec<Bson>) -> Self {
370        Self(doc! { key: { "$all": values } })
371    }
372
373    pub fn any(key: String, condition: Document) -> Self {
374        Self(doc! { key: { "$elemMatch": condition } })
375    }
376}
377
378impl From<Filter<serde_json::Value>> for MongoDbSearchFilter {
379    fn from(value: Filter<serde_json::Value>) -> Self {
380        value.interpret_with(|v| to_bson(&v).unwrap_or(Bson::Null))
381    }
382}
383
384impl DynamicSearchFilter for MongoDbSearchFilter {
385    fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError> {
386        Ok(filter.into())
387    }
388}
389
390impl<C, M> VectorStoreIndex for MongoDbVectorIndex<C, M>
391where
392    C: Sync + Send,
393    M: EmbeddingModel + Sync + Send,
394{
395    type Filter = MongoDbSearchFilter;
396
397    /// Implement the `top_n` method of the `VectorStoreIndex` trait for `MongoDbVectorIndex`.
398    ///
399    /// `VectorSearchRequest` similarity search threshold filter gets ignored here because it is already present and can already be added in the MongoDB vector store struct.
400    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
401        &self,
402        req: VectorSearchRequest<MongoDbSearchFilter>,
403    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
404        let project_stage = doc! {
405            "$project": {
406                self.embedded_field.clone(): 0
407            }
408        };
409
410        self.run_search_pipeline(&req, project_stage)
411            .await?
412            .into_iter()
413            .map(|(score, id, doc)| {
414                let doc_t: T = serde_json::from_value(doc).map_err(VectorStoreError::JsonError)?;
415                Ok((score, id, doc_t))
416            })
417            .collect()
418    }
419
420    /// Implement the `top_n_ids` method of the `VectorStoreIndex` trait for `MongoDbVectorIndex`.
421    async fn top_n_ids(
422        &self,
423        req: VectorSearchRequest<MongoDbSearchFilter>,
424    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
425        let project_stage = doc! {
426            "$project": {
427                "_id": 1,
428                "score": 1
429            },
430        };
431
432        Ok(self
433            .run_search_pipeline(&req, project_stage)
434            .await?
435            .into_iter()
436            .map(|(score, id, _)| (score, id))
437            .collect())
438    }
439}
440
441impl<C, M> InsertDocuments for MongoDbVectorIndex<C, M>
442where
443    C: Send + Sync,
444    M: EmbeddingModel + Send + Sync,
445{
446    async fn insert_documents<Doc: Serialize + Embed + Send>(
447        &self,
448        documents: Vec<(Doc, Vec<Embedding>)>,
449    ) -> Result<(), VectorStoreError> {
450        let mongo_documents = rig_core::vector_store::flatten_embedded(
451            documents,
452            |json_doc, embedding| {
453                Ok(doc! {
454                    "document": mongodb::bson::to_bson(json_doc).map_err(VectorStoreError::datastore)?,
455                    "embedding": embedding.vec,
456                    "embedded_text": embedding.document,
457                })
458            },
459        )?;
460
461        let collection = self.collection.clone_with_type::<mongodb::bson::Document>();
462
463        collection
464            .insert_many(mongo_documents)
465            .await
466            .map_err(VectorStoreError::datastore)?;
467
468        Ok(())
469    }
470}
rig_mongodb/lib.rs

rig_mongodb/
lib.rs
