feat: fastembed integration

Cargo.toml

@@ -9,5 +9,5 @@ members = [
     "rig-qdrant",
     "rig-core/rig-core-derive",
     "rig-sqlite",
-    "rig-eternalai",
+    "rig-eternalai", "rig-fastembed",
 ]

README.md

@@ -110,6 +110,9 @@ Vector stores are available as separate companion-crates:
 - Qdrant vector store: [`rig-qdrant`](https://github.com/0xPlaygrounds/rig/tree/main/rig-qdrant)
 - SQLite vector store: [`rig-sqlite`](https://github.com/0xPlaygrounds/rig/tree/main/rig-sqlite)
 
+The following providers are available as separate companion-crates:
+- Fastembed: [`rig-fastembed`](https://github.com/0xPlaygrounds/rig/tree/main/rig-fastembed)
+
 
 <p align="center">
 <br>

rig-core/examples/agent_evaluator_optimizer.rs

@@ -79,7 +79,7 @@ async fn main() -> Result<(), anyhow::Error> {
         } else {
             let context = format!("{TASK}\n\n{}", eval_result.feedback);
 
-            response = generator_agent.prompt(&context).await.unwrap();
+            response = generator_agent.prompt(context).await.unwrap();
             memories.push(response.clone());
         }
     }

rig-fastembed/.gitignore

@@ -0,0 +1 @@
+.fastembed_cache

rig-fastembed/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "rig-fastembed"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+rig-core = { path = "../rig-core", version = "0.7.0" }
+serde = { version = "1.0.193", features = ["derive"] }
+serde_json = "1.0.108"
+tracing = "0.1.40"
+schemars = "0.8.16"
+fastembed = "4.4.0"
+
+[dev-dependencies]
+anyhow = "1.0.75"
+tokio = { version = "1.34.0", features = ["full"] }
+
+[[example]]
+name = "vector_search"
+required-features = ["rig-core/derive"]

rig-fastembed/README.md

@@ -0,0 +1,4 @@
+## Fastembed integration with Rig
+This crate allows you to use [`fastembed-rs`](https://github.com/Anush008/fastembed-rs) with Rig.
+
+Unlike the providers found in the core crate, `fastembed` does not compile to the `wasm32-unknown-unknown` target.

rig-fastembed/examples/vector_search.rs

@@ -0,0 +1,81 @@
+use rig::{
+    embeddings::EmbeddingsBuilder,
+    vector_store::{in_memory_store::InMemoryVectorStore, VectorStoreIndex},
+    Embed,
+};
+use rig_fastembed::FastembedModel;
+use serde::{Deserialize, Serialize};
+
+// Shape of data that needs to be RAG'ed.
+// The definition field will be used to generate embeddings.
+#[derive(Embed, Clone, Deserialize, Debug, Serialize, Eq, PartialEq, Default)]
+struct WordDefinition {
+    id: String,
+    word: String,
+    #[embed]
+    definitions: Vec<String>,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    // Create OpenAI client
+    let fastembed_client = rig_fastembed::Client::new();
+
+    let embedding_model = fastembed_client.embedding_model(&FastembedModel::AllMiniLML6V2Q);
+
+    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
+        .documents(vec![
+            WordDefinition {
+                id: "doc0".to_string(),
+                word: "flurbo".to_string(),
+                definitions: vec![
+                    "A green alien that lives on cold planets.".to_string(),
+                    "A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
+                ]
+            },
+            WordDefinition {
+                id: "doc1".to_string(),
+                word: "glarb-glarb".to_string(),
+                definitions: vec![
+                    "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
+                    "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
+                ]
+            },
+            WordDefinition {
+                id: "doc2".to_string(),
+                word: "linglingdong".to_string(),
+                definitions: vec![
+                    "A term used by inhabitants of the sombrero galaxy to describe humans.".to_string(),
+                    "A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string()
+                ]
+            },
+        ])?
+        .build()
+        .await?;
+
+    // Create vector store with the embeddings
+    let vector_store =
+        InMemoryVectorStore::from_documents_with_id_f(embeddings, |doc| doc.id.clone());
+
+    // Create vector store index
+    let index = vector_store.index(embedding_model);
+
+    let results = index
+        .top_n::<WordDefinition>("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
+        .await?
+        .into_iter()
+        .map(|(score, id, doc)| (score, id, doc.word))
+        .collect::<Vec<_>>();
+
+    println!("Results: {:?}", results);
+
+    let id_results = index
+        .top_n_ids("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
+        .await?
+        .into_iter()
+        .collect::<Vec<_>>();
+
+    println!("ID results: {:?}", id_results);
+
+    Ok(())
+}

rig-fastembed/src/lib.rs

@@ -0,0 +1,154 @@
+use std::sync::Arc;
+
+pub use fastembed::EmbeddingModel as FastembedModel;
+use fastembed::{InitOptions, TextEmbedding};
+use rig::{
+    embeddings::{self, EmbeddingError, EmbeddingsBuilder},
+    Embed,
+};
+
+#[derive(Clone)]
+pub struct Client;
+
+impl Default for Client {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Client {
+    /// Create a new Fastembed client.
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Create an embedding model with the given name.
+    /// Note: default embedding dimension of 0 will be used if model is not known.
+    /// If this is the case, it's better to use function `embedding_model_with_ndims`
+    ///
+    /// # Example
+    /// ```
+    /// use rig_fastembed::{Client, FastembedModel};
+    ///
+    /// // Initialize the OpenAI client
+    /// let fastembed_client = Client::new("your-open-ai-api-key");
+    ///
+    /// let embedding_model = fastembed_client.embedding_model(&FastembedModel::AllMiniLML6V2Q);
+    /// ```
+    pub fn embedding_model(&self, model: &FastembedModel) -> EmbeddingModel {
+        let ndims = fetch_model_ndims(model);
+
+        EmbeddingModel::new(model, ndims)
+    }
+
+    /// Create an embedding builder with the given embedding model.
+    ///
+    /// # Example
+    /// ```
+    /// use rig_fastembed::{Client, FastembedModel};
+    ///
+    /// // Initialize the Fastembed client
+    /// let fastembed_client = Client::new();
+    ///
+    /// let embeddings = fastembed_client.embeddings(FastembedModel::AllMiniLML6V2Q)
+    ///     .simple_document("doc0", "Hello, world!")
+    ///     .simple_document("doc1", "Goodbye, world!")
+    ///     .build()
+    ///     .await
+    ///     .expect("Failed to embed documents");
+    /// ```
+    pub fn embeddings<D: Embed>(
+        &self,
+        model: &fastembed::EmbeddingModel,
+    ) -> EmbeddingsBuilder<EmbeddingModel, D> {
+        EmbeddingsBuilder::new(self.embedding_model(model))
+    }
+}
+
+#[derive(Clone)]
+pub struct EmbeddingModel {
+    embedder: Arc<TextEmbedding>,
+    pub model: FastembedModel,
+    ndims: usize,
+}
+
+impl EmbeddingModel {
+    pub fn new(model: &fastembed::EmbeddingModel, ndims: usize) -> Self {
+        let embedder = Arc::new(
+            TextEmbedding::try_new(
+                InitOptions::new(model.to_owned()).with_show_download_progress(true),
+            )
+            .unwrap(),
+        );
+
+        Self {
+            embedder,
+            model: model.to_owned(),
+            ndims,
+        }
+    }
+}
+
+impl embeddings::EmbeddingModel for EmbeddingModel {
+    const MAX_DOCUMENTS: usize = 1024;
+
+    fn ndims(&self) -> usize {
+        self.ndims
+    }
+
+    async fn embed_texts(
+        &self,
+        documents: impl IntoIterator<Item = String>,
+    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
+        let documents_as_strings: Vec<String> = documents.into_iter().collect();
+
+        let documents_as_vec = self
+            .embedder
+            .embed(documents_as_strings.clone(), None)
+            .map_err(|err| EmbeddingError::ProviderError(err.to_string()))?;
+
+        let docs = documents_as_strings
+            .into_iter()
+            .zip(documents_as_vec)
+            .map(|(document, embedding)| embeddings::Embedding {
+                document,
+                vec: embedding.into_iter().map(|f| f as f64).collect(),
+            })
+            .collect::<Vec<embeddings::Embedding>>();
+
+        Ok(docs)
+    }
+}
+
+/// As seen on the text embedding model cards file: <https://github.com/Anush008/fastembed-rs/blob/main/src/models/text_embedding.rs>
+pub fn fetch_model_ndims(model: &FastembedModel) -> usize {
+    match model {
+        FastembedModel::AllMiniLML6V2
+        | FastembedModel::AllMiniLML6V2Q
+        | FastembedModel::AllMiniLML12V2
+        | FastembedModel::AllMiniLML12V2Q
+        | FastembedModel::BGESmallENV15
+        | FastembedModel::BGESmallENV15Q
+        | FastembedModel::ParaphraseMLMiniLML12V2Q
+        | FastembedModel::ParaphraseMLMiniLML12V2
+        | FastembedModel::MultilingualE5Small => 384,
+        FastembedModel::BGESmallZHV15 | FastembedModel::ClipVitB32 => 512,
+        FastembedModel::BGEBaseENV15
+        | FastembedModel::BGEBaseENV15Q
+        | FastembedModel::NomicEmbedTextV1
+        | FastembedModel::NomicEmbedTextV15
+        | FastembedModel::NomicEmbedTextV15Q
+        | FastembedModel::ParaphraseMLMpnetBaseV2
+        | FastembedModel::MultilingualE5Base
+        | FastembedModel::GTEBaseENV15
+        | FastembedModel::GTEBaseENV15Q
+        | FastembedModel::JinaEmbeddingsV2BaseCode => 768,
+        FastembedModel::BGELargeENV15
+        | FastembedModel::BGELargeENV15Q
+        | FastembedModel::MultilingualE5Large
+        | FastembedModel::MxbaiEmbedLargeV1
+        | FastembedModel::MxbaiEmbedLargeV1Q
+        | FastembedModel::GTELargeENV15
+        | FastembedModel::GTELargeENV15Q => 1024,
+    }
+}