fix(mongodb): remove embeddings from top_n lookup

rig-core/src/vector_store/mod.rs

@@ -17,6 +17,9 @@ pub enum VectorStoreError {
 
     #[error("Datastore error: {0}")]
     DatastoreError(#[from] Box<dyn std::error::Error + Send + Sync>),
+
+    #[error("Vector store error: {0}")]
+    Error(String),
 }
 
 /// Trait for vector stores

rig-mongodb/examples/vector_search_mongodb.rs

@@ -5,7 +5,7 @@ use rig::{
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     vector_store::VectorStoreIndex,
 };
-use rig_mongodb::{MongoDbVectorStore, SearchParams};
+use rig_mongodb::{DocumentResponse, MongoDbVectorStore, SearchParams};
 use std::env;
 
 #[tokio::main]
@@ -49,11 +49,13 @@ async fn main() -> Result<(), anyhow::Error> {
 
     // Create a vector index on our vector store
     // IMPORTANT: Reuse the same model that was used to generate the embeddings
-    let index = vector_store.index(model, "vector_index", SearchParams::default());
+    let index = vector_store
+        .index(model, "vector_index", SearchParams::default())
+        .await?;
 
     // Query the index
     let results = index
-        .top_n::<DocumentEmbeddings>("What is a linglingdong?", 1)
+        .top_n::<DocumentResponse>("What is a linglingdong?", 1)
         .await?
         .into_iter()
         .map(|(score, id, doc)| (score, id, doc.document))

rig-mongodb/src/lib.rs

@@ -5,13 +5,58 @@ use rig::{
     embeddings::{DocumentEmbeddings, Embedding, EmbeddingModel},
     vector_store::{VectorStore, VectorStoreError, VectorStoreIndex},
 };
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 
 /// A MongoDB vector store.
 pub struct MongoDbVectorStore {
     collection: mongodb::Collection<DocumentEmbeddings>,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct SearchIndex {
+    id: String,
+    name: String,
+    #[serde(rename = "type")]
+    index_type: String,
+    status: String,
+    queryable: bool,
+    latest_definition: LatestDefinition,
+}
+
+impl SearchIndex {
+    async fn get_search_index(
+        collection: mongodb::Collection<DocumentEmbeddings>,
+        index_name: &str,
+    ) -> Result<SearchIndex, VectorStoreError> {
+        collection
+            .list_search_indexes(index_name, None, None)
+            .await
+            .map_err(mongodb_to_rig_error)?
+            .with_type::<SearchIndex>()
+            .next()
+            .await
+            .transpose()
+            .map_err(mongodb_to_rig_error)?
+            .ok_or(VectorStoreError::Error("Index not found".to_string()))
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct LatestDefinition {
+    fields: Vec<Field>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct Field {
+    #[serde(rename = "type")]
+    field_type: String,
+    path: String,
+    num_dimensions: i32,
+    similarity: String,
+}
+
 fn mongodb_to_rig_error(e: mongodb::error::Error) -> VectorStoreError {
     VectorStoreError::DatastoreError(Box::new(e))
 }
@@ -87,13 +132,13 @@ impl MongoDbVectorStore {
     ///
     /// The index (of type "vector") must already exist for the MongoDB collection.
     /// See the MongoDB [documentation](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/) for more information on creating indexes.
-    pub fn index<M: EmbeddingModel>(
+    pub async fn index<M: EmbeddingModel>(
         &self,
         model: M,
         index_name: &str,
         search_params: SearchParams,
-    ) -> MongoDbVectorIndex<M> {
-        MongoDbVectorIndex::new(self.collection.clone(), model, index_name, search_params)
+    ) -> Result<MongoDbVectorIndex<M>, VectorStoreError> {
+        MongoDbVectorIndex::new(self.collection.clone(), model, index_name, search_params).await
     }
 }
 
@@ -102,6 +147,7 @@ pub struct MongoDbVectorIndex<M: EmbeddingModel> {
     collection: mongodb::Collection<DocumentEmbeddings>,
     model: M,
     index_name: String,
+    embedded_field: String,
     search_params: SearchParams,
 }
 
@@ -118,7 +164,7 @@ impl<M: EmbeddingModel> MongoDbVectorIndex<M> {
         doc! {
           "$vectorSearch": {
             "index": &self.index_name,
-            "path": "embeddings.vec",
+            "path": self.embedded_field.clone(),
             "queryVector": &prompt_embedding.vec,
             "numCandidates": num_candidates.unwrap_or((n * 10) as u32),
             "limit": n as u32,
@@ -140,22 +186,42 @@ impl<M: EmbeddingModel> MongoDbVectorIndex<M> {
 }
 
 impl<M: EmbeddingModel> MongoDbVectorIndex<M> {
-    pub fn new(
+    pub async fn new(
         collection: mongodb::Collection<DocumentEmbeddings>,
         model: M,
         index_name: &str,
         search_params: SearchParams,
-    ) -> Self {
-        Self {
+    ) -> Result<Self, VectorStoreError> {
+        let search_index = SearchIndex::get_search_index(collection.clone(), index_name).await?;
+
+        if !search_index.queryable {
+            return Err(VectorStoreError::Error(
+                "Index is not queryable".to_string(),
+            ));
+        }
+
+        let embedded_field = search_index
+            .latest_definition
+            .fields
+            .into_iter()
+            .map(|field| field.path)
+            .next()
+            // This error shouldn't occur if the index is queryable
+            .ok_or(VectorStoreError::Error(
+                "No embedded fields found".to_string(),
+            ))?;
+
+        Ok(Self {
             collection,
             model,
             index_name: index_name.to_string(),
+            embedded_field,
             search_params,
-        }
+        })
     }
 }
 
-/// See [MongoDB Vector Search](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/) for more information
+/// See [MongoDB Vector Search](`https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/`) for more information
 /// on each of the fields
 pub struct SearchParams {
     filter: mongodb::bson::Document,
@@ -219,6 +285,13 @@ impl<M: EmbeddingModel + std::marker::Sync + Send> VectorStoreIndex for MongoDbV
                 [
                     self.pipeline_search_stage(&prompt_embedding, n),
                     self.pipeline_score_stage(),
+                    {
+                        doc! {
+                            "$project": {
+                                self.embedded_field.clone(): 0,
+                            },
+                        }
+                    },
                 ],
                 None,
             )
@@ -291,3 +364,10 @@ impl<M: EmbeddingModel + std::marker::Sync + Send> VectorStoreIndex for MongoDbV
         Ok(results)
     }
 }
+
+#[derive(Clone, Eq, PartialEq, Serialize, Deserialize, Debug)]
+pub struct DocumentResponse {
+    #[serde(rename = "_id")]
+    pub id: String,
+    pub document: serde_json::Value,
+}