feat(rig-841): support HelixDB (#805)

joshua-mo-143 · web-flow · commit e6f84c680691 · 2025-10-06T17:46:19.000+01:00
* feat(rig-841): support helixdb

* docs: license, update cargo.toml, readme

* docs: minor docs update

* feat: impl InsertDocuments for HelixDBVectorStore

* refactor: update docs

* chore: update deps

* refactor: fix stuff

* docs: update docs, licensing

* refactor: add docstrings
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,7 +16,7 @@ members = [
     "rig-bedrock",
     "rig-milvus",
     "rig-wasm",
-    "rig-s3vectors",
+    "rig-s3vectors", "rig-helixdb",
 ]
 
 [profile.release]
diff --git a/README.md b/README.md
@@ -122,6 +122,7 @@ Vector stores are available as separate companion-crates:
 - Milvus: [`rig-milvus`](https://github.com/0xPlaygrounds/rig/tree/main/rig-milvus)
 - ScyllaDB: [`rig-scylladb`](https://github.com/0xPlaygrounds/rig/tree/main/rig-scylladb)
 - AWS S3Vectors: [`rig-s3vectors`](https://github.com/0xPlaygrounds/rig/tree/main/rig-s3vectors)
+- HelixDB: [`rig-helixdb`](https://github.com/0xPlaygrounds/rig/tree/main/rig-helixdb)
 
 The following providers are available as separate companion-crates:
 - Fastembed: [`rig-fastembed`](https://github.com/0xPlaygrounds/rig/tree/main/rig-fastembed)
diff --git a/rig-helixdb/.gitignore b/rig-helixdb/.gitignore
@@ -0,0 +1 @@
+.helix
diff --git a/rig-helixdb/Cargo.toml b/rig-helixdb/Cargo.toml
@@ -0,0 +1,21 @@
+[package]
+name = "rig-helixdb"
+version = "0.1.0"
+edition.workspace = true
+license = "MIT"
+readme = "README.md"
+description = "Rig vector store index integration for HelixDB."
+repository = "https://github.com/0xPlaygrounds/rig"
+
+[dependencies]
+helix-rs = "0.1.9"
+serde = { workspace = true, features = ["derive"] }
+serde_json.workspace = true
+rig-core = { path = "../rig-core", version = "0.21.0" }
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+
+[[example]]
+name = "vector_search_helixdb"
+required-features = ["rig-core/derive"]
diff --git a/rig-helixdb/README.md b/rig-helixdb/README.md
@@ -0,0 +1,35 @@
+# Rig HelixDB integration
+This crate integrates HelixDB into Rig, allowing you to easily use RAG with this database.
+
+## Installation
+To install this crate, run the following command in a Rust project directory which will add `rig-helixdb` as a dependency (requires `rig-core` added for intended usage):
+```bash
+cargo add rig-helixdb
+```
+
+There's a few different ways you can run HelixDB:
+- Through HelixDB's cloud offering
+- Running it locally through their `helix start` command (requires the Helix CLI to be installed).
+  - For local dev, you will likely want to use `helix push dev` for continuous iteration.
+
+## How to run the example
+Before running the example, you will need to ensure that you are running an instance of HelixDB which you can do with `helix dockerdev run`.
+
+Once done, you will then need to deploy your queries/schema. **The queries/schema in the `examples/helixdb-cfg` folder are a required minimum to be use this integration.** `rig-helixdb` also additionally provides a way to get a manual handle on the client yourself so that you can invoke your own queries should you need to.
+
+Assuming `rig-helixdb` is your current working directory, to deploy a minimum viable configuration for HelixDB (with `rig-helixdb`) you will need to `cd` into the `helixdb-cfg` folder and then run the following:
+```bash
+helix push dev
+```
+
+This will then deploy the queries/schema into your instance.
+
+To run the example, add your OpenAI API key as an environment variable:
+```bash
+export OPENAI_API_KEY=my_key
+```
+
+Finally, use the following command below to run the example:
+```bash
+cargo run --example vector_search_helixdb --features rig-core/derive
+```
diff --git a/rig-helixdb/examples/helixdb-cfg/db/queries.hx b/rig-helixdb/examples/helixdb-cfg/db/queries.hx
@@ -0,0 +1,7 @@
+QUERY InsertVector (vector: [F64], doc: String, json_payload: String) =>
+    AddV<Document>(vector, { doc: doc, json_payload: json_payload })
+    RETURN doc
+
+QUERY VectorSearch(vector: [F64], limit: U64, threshold: F64) =>
+    vec_docs <- SearchV<Document>(vector, limit)
+    RETURN vec_docs
diff --git a/rig-helixdb/examples/helixdb-cfg/db/schema.hx b/rig-helixdb/examples/helixdb-cfg/db/schema.hx
@@ -0,0 +1,4 @@
+V::Document {
+    doc: String,
+    json_payload: String
+}
diff --git a/rig-helixdb/examples/helixdb-cfg/helix.toml b/rig-helixdb/examples/helixdb-cfg/helix.toml
@@ -0,0 +1,15 @@
+[project]
+name = "helixdb-cfg"
+queries = "./db/"
+
+[local.dev]
+port = 6969
+build_mode = "debug"
+
+[local.dev.vector_config]
+m = 16
+ef_construction = 128
+ef_search = 768
+db_max_size_gb = 10
+
+[cloud]
diff --git a/rig-helixdb/examples/vector_search_helixdb.rs b/rig-helixdb/examples/vector_search_helixdb.rs
@@ -0,0 +1,79 @@
+use helix_rs::{HelixDB, HelixDBClient};
+use rig::{
+    Embed,
+    client::{EmbeddingsClient, ProviderClient},
+    embeddings::EmbeddingsBuilder,
+    vector_store::{InsertDocuments, VectorSearchRequest, VectorStoreIndex},
+};
+use rig_helixdb::HelixDBVectorStore;
+use serde::{Deserialize, Serialize};
+
+// A vector search needs to be performed on the `definitions` field, so we derive the `Embed` trait for `WordDefinition`
+// and tag that field with `#[embed]`.
+// We are not going to store the definitions on our database so we skip the `Serialize` trait
+#[derive(Embed, Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Default)]
+struct WordDefinition {
+    word: String,
+    #[serde(skip)] // we don't want to serialize this field, we use only to create embeddings
+    #[embed]
+    definition: String,
+}
+
+impl std::fmt::Display for WordDefinition {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.word)
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let openai_model =
+        rig::providers::openai::Client::from_env().embedding_model("text-embedding-ada-002");
+
+    let helixdb_client = HelixDB::new(None, Some(6969), None); // Uses default port 6969
+    let vector_store = HelixDBVectorStore::new(helixdb_client, openai_model.clone());
+
+    let words = vec![
+        WordDefinition {
+            word: "flurbo".to_string(),
+            definition: "1. *flurbo* (name): A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
+        },
+        WordDefinition {
+            word: "glarb-glarb".to_string(),
+            definition: "1. *glarb-glarb* (noun): A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
+        },
+        WordDefinition {
+            word: "linglingdong".to_string(),
+            definition: "1. *linglingdong* (noun): A term used by inhabitants of the far side of the moon to describe humans.".to_string(),
+        }];
+
+    let documents = EmbeddingsBuilder::new(openai_model)
+        .documents(words)
+        .unwrap()
+        .build()
+        .await
+        .expect("Failed to create embeddings");
+
+    vector_store.insert_documents(documents).await.unwrap();
+
+    let query = "What is a flurbo?";
+    let vector_req = VectorSearchRequest::builder()
+        .query(query)
+        .samples(5)
+        .build()
+        .unwrap();
+
+    let docs = vector_store
+        .top_n::<WordDefinition>(vector_req)
+        .await
+        .unwrap();
+
+    for doc in docs {
+        println!(
+            "Vector found with id: {id} and score: {score} and word def: {doc}",
+            id = doc.1,
+            score = doc.0,
+            doc = doc.2
+        )
+    }
+}
diff --git a/rig-helixdb/src/lib.rs b/rig-helixdb/src/lib.rs
@@ -0,0 +1,181 @@
+use helix_rs::HelixDBClient;
+use rig::{
+    embeddings::EmbeddingModel,
+    vector_store::{InsertDocuments, VectorStoreError, VectorStoreIndex},
+};
+use serde::{Deserialize, Serialize};
+
+/// A client for easily carrying out Rig-related vector store operations.
+///
+/// If you are unsure what type to use for the client, `helix_rs::HelixDB` is the typical default.
+///
+/// Usage:
+/// ```rust
+/// let openai_model =
+///     rig::providers::openai::Client::from_env().embedding_model("text-embedding-ada-002");
+///
+/// let helixdb_client = HelixDB::new(None, Some(6969), None);
+/// let vector_store = HelixDBVectorStore::new(helixdb_client, openai_model.clone());
+/// ```
+pub struct HelixDBVectorStore<C, E> {
+    client: C,
+    model: E,
+}
+
+/// The result of a query. Only used internally as this is a representative type required for the relevant HelixDB query (`VectorSearch`).
+#[derive(Deserialize, Serialize, Clone, Debug)]
+struct QueryResult {
+    id: String,
+    score: f64,
+    doc: String,
+    json_payload: String,
+}
+
+/// An input query. Only used internally as this is a representative type required for the relevant HelixDB query (`VectorSearch`).
+#[derive(Deserialize, Serialize, Clone, Debug)]
+struct QueryInput {
+    vector: Vec<f64>,
+    limit: u64,
+    threshold: f64,
+}
+
+impl QueryInput {
+    /// Makes a new instance of `QueryInput`.
+    pub(crate) fn new(vector: Vec<f64>, limit: u64, threshold: f64) -> Self {
+        Self {
+            vector,
+            limit,
+            threshold,
+        }
+    }
+}
+
+impl<C, E> HelixDBVectorStore<C, E>
+where
+    C: HelixDBClient + Send,
+    E: EmbeddingModel,
+{
+    pub fn new(client: C, model: E) -> Self {
+        Self { client, model }
+    }
+
+    pub fn client(&self) -> &C {
+        &self.client
+    }
+}
+
+impl<C, E> InsertDocuments for HelixDBVectorStore<C, E>
+where
+    C: HelixDBClient + Send + Sync,
+    E: EmbeddingModel + Send + Sync,
+{
+    async fn insert_documents<Doc: Serialize + rig::Embed + Send>(
+        &self,
+        documents: Vec<(Doc, rig::OneOrMany<rig::embeddings::Embedding>)>,
+    ) -> Result<(), VectorStoreError> {
+        #[derive(Serialize, Deserialize, Clone, Debug, Default)]
+        struct QueryInput {
+            vector: Vec<f64>,
+            doc: String,
+            json_payload: String,
+        }
+
+        #[derive(Serialize, Deserialize, Clone, Debug, Default)]
+        struct QueryOutput {
+            doc: String,
+        }
+
+        for (document, embeddings) in documents {
+            let json_document: serde_json::Value = serde_json::to_value(&document).unwrap();
+            let json_document_as_string = serde_json::to_string(&json_document).unwrap();
+
+            for embedding in embeddings {
+                let embedded_text = embedding.document;
+                let vector: Vec<f64> = embedding.vec;
+
+                let query = QueryInput {
+                    vector,
+                    doc: embedded_text,
+                    json_payload: json_document_as_string.clone(),
+                };
+
+                self.client
+                    .query::<QueryInput, QueryOutput>("InsertVector", &query)
+                    .await
+                    .inspect_err(|x| println!("Error: {x}"))
+                    .map_err(|x| VectorStoreError::DatastoreError(x.to_string().into()))?;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl<C, E> VectorStoreIndex for HelixDBVectorStore<C, E>
+where
+    C: HelixDBClient + Send + Sync,
+    E: EmbeddingModel + Send + Sync,
+{
+    async fn top_n<T: for<'a> serde::Deserialize<'a> + Send>(
+        &self,
+        req: rig::vector_store::VectorSearchRequest,
+    ) -> Result<Vec<(f64, String, T)>, rig::vector_store::VectorStoreError> {
+        let vector = self.model.embed_text(req.query()).await?.vec;
+
+        let query_input =
+            QueryInput::new(vector, req.samples(), req.threshold().unwrap_or_default());
+
+        #[derive(Serialize, Deserialize, Debug)]
+        struct VecResult {
+            vec_docs: Vec<QueryResult>,
+        }
+
+        let result: VecResult = self
+            .client
+            .query::<QueryInput, VecResult>("VectorSearch", &query_input)
+            .await
+            .unwrap();
+
+        let docs = result
+            .vec_docs
+            .into_iter()
+            .map(|x| {
+                let doc: T = serde_json::from_str(&x.json_payload)?;
+
+                // HelixDB gives us the cosine distance, so we need to use `-(cosine_dist - 1)` to get the cosine similarity score.
+                Ok((-(x.score - 1.), x.id, doc))
+            })
+            .collect::<Result<Vec<_>, VectorStoreError>>()?;
+
+        Ok(docs)
+    }
+
+    async fn top_n_ids(
+        &self,
+        req: rig::vector_store::VectorSearchRequest,
+    ) -> Result<Vec<(f64, String)>, rig::vector_store::VectorStoreError> {
+        let vector = self.model.embed_text(req.query()).await?.vec;
+
+        let query_input =
+            QueryInput::new(vector, req.samples(), req.threshold().unwrap_or_default());
+
+        #[derive(Serialize, Deserialize, Debug)]
+        struct VecResult {
+            vec_docs: Vec<QueryResult>,
+        }
+
+        let result: VecResult = self
+            .client
+            .query::<QueryInput, VecResult>("VectorSearch", &query_input)
+            .await
+            .unwrap();
+
+        // HelixDB gives us the cosine distance, so we need to use `-(cosine_dist - 1)` to get the cosine similarity score.
+        let docs = result
+            .vec_docs
+            .into_iter()
+            .map(|x| Ok((-(x.score - 1.), x.id)))
+            .collect::<Result<Vec<_>, VectorStoreError>>()?;
+
+        Ok(docs)
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ members = [`
`16`	`16`	`"rig-bedrock",`
`17`	`17`	`"rig-milvus",`
`18`	`18`	`"rig-wasm",`
`19`		`- "rig-s3vectors",`
	`19`	`+ "rig-s3vectors", "rig-helixdb",`
`20`	`20`	`]`
`21`	`21`
`22`	`22`	`[profile.release]`