refactor: remove DocumentEmbeddings from in memory vector store

0xPlaygrounds · Oct 9, 2024 · 5407772 · 5407772
1 parent f5441db
commit 5407772
Show file tree

Hide file tree

Showing 9 changed files with 134 additions and 270 deletions.
diff --git a/rig-core/examples/calculator_chatbot.rs b/rig-core/examples/calculator_chatbot.rs
@@ -2,10 +2,10 @@ use anyhow::Result;
 use rig::{
     cli_chatbot::cli_chatbot,
     completion::ToolDefinition,
-    embeddings::EmbeddingsBuilder,
+    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     tool::{Tool, ToolEmbedding, ToolSet},
-    vector_store::{in_memory_store::InMemoryVectorStore, VectorStore},
+    vector_store::in_memory_store::InMemoryVectorStore,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;
@@ -251,9 +251,19 @@ async fn main() -> Result<(), anyhow::Error> {
         .build()
         .await?;
 
-    let mut store = InMemoryVectorStore::default();
-    store.add_documents(embeddings).await?;
-    let index = store.index(embedding_model);
+    let vector_store = InMemoryVectorStore::default().add_documents(
+        embeddings
+            .into_iter()
+            .map(
+                |DocumentEmbeddings {
+                     id,
+                     document,
+                     embeddings,
+                 }| { (id, document, embeddings) },
+            )
+            .collect(),
+    )?;
+    let index = vector_store.index(embedding_model);
 
     // Create RAG agent with a single context prompt and a dynamic tool source
     let calculator_rag = openai_client

diff --git a/rig-core/examples/rag.rs b/rig-core/examples/rag.rs
@@ -2,9 +2,9 @@ use std::env;
 
 use rig::{
     completion::Prompt,
-    embeddings::EmbeddingsBuilder,
+    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
-    vector_store::{in_memory_store::InMemoryVectorStore, VectorStore},
+    vector_store::in_memory_store::InMemoryVectorStore,
 };
 
 #[tokio::main]
@@ -25,7 +25,18 @@ async fn main() -> Result<(), anyhow::Error> {
         .build()
         .await?;
 
-    vector_store.add_documents(embeddings).await?;
+    let vector_store = vector_store.add_documents(
+        embeddings
+            .into_iter()
+            .map(
+                |DocumentEmbeddings {
+                     id,
+                     document,
+                     embeddings,
+                 }| { (id, document, embeddings) },
+            )
+            .collect(),
+    )?;
 
     // Create vector store index
     let index = vector_store.index(embedding_model);

diff --git a/rig-core/examples/rag_dynamic_tools.rs b/rig-core/examples/rag_dynamic_tools.rs
@@ -1,10 +1,10 @@
 use anyhow::Result;
 use rig::{
     completion::{Prompt, ToolDefinition},
-    embeddings::EmbeddingsBuilder,
+    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     tool::{Tool, ToolEmbedding, ToolSet},
-    vector_store::{in_memory_store::InMemoryVectorStore, VectorStore},
+    vector_store::in_memory_store::InMemoryVectorStore,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;
@@ -150,9 +150,6 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let embedding_model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);
 
-    // Create vector store, compute tool embeddings and load them in the store
-    let mut vector_store = InMemoryVectorStore::default();
-
     let toolset = ToolSet::builder()
         .dynamic_tool(Add)
         .dynamic_tool(Subtract)
@@ -163,7 +160,18 @@ async fn main() -> Result<(), anyhow::Error> {
         .build()
         .await?;
 
-    vector_store.add_documents(embeddings).await?;
+    let vector_store = InMemoryVectorStore::default().add_documents(
+        embeddings
+            .into_iter()
+            .map(
+                |DocumentEmbeddings {
+                     id,
+                     document,
+                     embeddings,
+                 }| { (id, document, embeddings) },
+            )
+            .collect(),
+    )?;
 
     // Create vector store index
     let index = vector_store.index(embedding_model);

diff --git a/rig-core/examples/vector_search.rs b/rig-core/examples/vector_search.rs
@@ -3,7 +3,10 @@ use std::env;
 use rig::{
     embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
-    vector_store::{in_memory_store::InMemoryVectorIndex, VectorStoreIndex},
+    vector_store::{
+        in_memory_store::{InMemoryVectorIndex, InMemoryVectorStore},
+        VectorStoreIndex,
+    },
 };
 
 #[tokio::main]
@@ -21,7 +24,20 @@ async fn main() -> Result<(), anyhow::Error> {
         .build()
         .await?;
 
-    let index = InMemoryVectorIndex::from_embeddings(model, embeddings).await?;
+    let vector_store = InMemoryVectorStore::default().add_documents(
+        embeddings
+            .into_iter()
+            .map(
+                |DocumentEmbeddings {
+                     id,
+                     document,
+                     embeddings,
+                 }| { (id, document, embeddings) },
+            )
+            .collect(),
+    )?;
+
+    let index = vector_store.index(model);
 
     let results = index
         .top_n::<DocumentEmbeddings>("What is a linglingdong?", 1)

diff --git a/rig-core/examples/vector_search_cohere.rs b/rig-core/examples/vector_search_cohere.rs
@@ -3,7 +3,7 @@ use std::env;
 use rig::{
     embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
     providers::cohere::{Client, EMBED_ENGLISH_V3},
-    vector_store::{in_memory_store::InMemoryVectorStore, VectorStore, VectorStoreIndex},
+    vector_store::{in_memory_store::InMemoryVectorStore, VectorStoreIndex},
 };
 
 #[tokio::main]
@@ -15,16 +15,25 @@ async fn main() -> Result<(), anyhow::Error> {
     let document_model = cohere_client.embedding_model(EMBED_ENGLISH_V3, "search_document");
     let search_model = cohere_client.embedding_model(EMBED_ENGLISH_V3, "search_query");
 
-    let mut vector_store = InMemoryVectorStore::default();
-
     let embeddings = EmbeddingsBuilder::new(document_model)
         .simple_document("doc0", "Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets")
         .simple_document("doc1", "Definition of a *glarb-glarb*: A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.")
         .simple_document("doc2", "Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.")
         .build()
         .await?;
 
-    vector_store.add_documents(embeddings).await?;
+    let vector_store = InMemoryVectorStore::default().add_documents(
+        embeddings
+            .into_iter()
+            .map(
+                |DocumentEmbeddings {
+                     id,
+                     document,
+                     embeddings,
+                 }| { (id, document, embeddings) },
+            )
+            .collect(),
+    )?;
 
     let index = vector_store.index(search_model);