Merge pull request #81 from Anush008/qdrant

feat: Qdrant support

Cargo.toml

@@ -1,6 +1,3 @@
 [workspace]
 resolver = "2"
-members = [
-    "rig-core", "rig-lancedb",
-    "rig-mongodb",
-]
+members = ["rig-core", "rig-lancedb", "rig-mongodb", "rig-qdrant"]

rig-core/examples/vector_search.rs

@@ -36,7 +36,6 @@ async fn main() -> Result<(), anyhow::Error> {
         .top_n_ids("What is a linglingdong?", 1)
         .await?
         .into_iter()
-        .map(|(score, id)| (score, id))
         .collect::<Vec<_>>();
 
     println!("ID results: {:?}", id_results);

rig-mongodb/examples/vector_search_mongodb.rs

@@ -65,7 +65,6 @@ async fn main() -> Result<(), anyhow::Error> {
         .top_n_ids("What is a linglingdong?", 1)
         .await?
         .into_iter()
-        .map(|(score, id)| (score, id))
         .collect::<Vec<_>>();
 
     println!("ID results: {:?}", id_results);

rig-qdrant/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "rig-qdrant"
+version = "0.1.0"
+edition = "2021"
+license = "MIT"
+readme = "README.md"
+description = "Rig vector store index integration for Qdrant. https://qdrant.tech"
+repository = "https://github.com/0xPlaygrounds/rig"
+
+[dependencies]
+rig-core = { path = "../rig-core", version = "0.3.0" }
+serde_json = "1.0.128"
+serde = "1.0.210"
+qdrant-client = "1.12.1"
+
+[dev-dependencies]
+tokio = { version = "1.40.0", features = ["rt-multi-thread"] }
+anyhow = "1.0.89"

rig-qdrant/LICENSE

@@ -0,0 +1,7 @@
+Copyright (c) 2024, Playgrounds Analytics Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

rig-qdrant/README.md

@@ -0,0 +1,5 @@
+# Rig-Qdrant
+
+Vector store index integration for [Qdrant](https://qdrant.tech/). This integration supports dense vector retrieval using Rig's embedding providers. It is also extensible to allow all [hybrid queries](https://qdrant.tech/documentation/concepts/hybrid-queries/) supported by Qdrant.
+
+You can find end-to-end examples [here](https://github.com/0xPlaygrounds/rig/tree/main/rig-qdrant/examples).

rig-qdrant/examples/qdrant_vector_search.rs

@@ -0,0 +1,85 @@
+// To run this example:
+//
+// export OPENAI_API_KEY=<YOUR-API-KEY>
+// docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant
+// cargo run --release --example qdrant_vector_search
+//
+// You can view the data at http://localhost:6333/dashboard
+
+use std::env;
+
+use qdrant_client::{
+    qdrant::{
+        CreateCollectionBuilder, Distance, PointStruct, QueryPointsBuilder, UpsertPointsBuilder,
+        VectorParamsBuilder,
+    },
+    Payload, Qdrant,
+};
+use rig::{
+    embeddings::EmbeddingsBuilder,
+    providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
+    vector_store::VectorStoreIndex,
+};
+use rig_qdrant::QdrantVectorStore;
+use serde_json::json;
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    const COLLECTION_NAME: &str = "rig-collection";
+
+    let client = Qdrant::from_url("http://localhost:6334").build()?;
+
+    // Create a collection with 1536 dimensions if it doesn't exist
+    if !client.collection_exists(COLLECTION_NAME).await? {
+        client
+            .create_collection(
+                CreateCollectionBuilder::new(COLLECTION_NAME)
+                    .vectors_config(VectorParamsBuilder::new(1536, Distance::Cosine)),
+            )
+            .await?;
+    }
+
+    // Initialize OpenAI client.
+    // Get your API key from https://platform.openai.com/api-keys
+    let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
+    let openai_client = Client::new(&openai_api_key);
+
+    let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);
+
+    let documents = EmbeddingsBuilder::new(model.clone())
+        .simple_document("0981d983-a5f8-49eb-89ea-f7d3b2196d2e", "Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets")
+        .simple_document("62a36d43-80b6-4fd6-990c-f75bb02287d1", "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("f9e17d59-32e5-440c-be02-b2759a654824", "Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.")
+        .build()
+        .await?;
+
+    let points: Vec<PointStruct> = documents
+        .into_iter()
+        .map(|d| {
+            let vec: Vec<f32> = d.embeddings[0].vec.iter().map(|&x| x as f32).collect();
+            PointStruct::new(
+                d.id,
+                vec,
+                Payload::try_from(json!({
+                    "document": d.document,
+                }))
+                .unwrap(),
+            )
+        })
+        .collect();
+
+    client
+        .upsert_points(UpsertPointsBuilder::new(COLLECTION_NAME, points))
+        .await?;
+
+    let query_params = QueryPointsBuilder::new(COLLECTION_NAME).with_payload(true);
+    let vector_store = QdrantVectorStore::new(client, model, query_params.build());
+
+    let results = vector_store
+        .top_n::<serde_json::Value>("What is a linglingdong?", 1)
+        .await?;
+
+    println!("Results: {:?}", results);
+
+    Ok(())
+}

rig-qdrant/src/lib.rs

@@ -0,0 +1,123 @@
+use qdrant_client::{
+    qdrant::{point_id::PointIdOptions, PointId, Query, QueryPoints},
+    Qdrant,
+};
+use rig::{
+    embeddings::EmbeddingModel,
+    vector_store::{VectorStoreError, VectorStoreIndex},
+};
+use serde::Deserialize;
+
+/// Represents a vector store implementation using Qdrant - <https://qdrant.tech/> as the backend.
+pub struct QdrantVectorStore<M: EmbeddingModel> {
+    /// Model used to generate embeddings for the vector store
+    model: M,
+    /// Client instance for Qdrant server communication
+    client: Qdrant,
+    /// Default search parameters
+    query_params: QueryPoints,
+}
+
+impl<M: EmbeddingModel> QdrantVectorStore<M> {
+    /// Creates a new instance of QdrantVectorStore.
+    ///
+    /// # Arguments
+    /// * `client` - Qdrant client instance
+    /// * `model` - Embedding model instance
+    /// * `query_params` - Search parameters for vector queries
+    ///     Reference: <https://api.qdrant.tech/v-1-12-x/api-reference/search/query-points>
+    pub fn new(client: Qdrant, model: M, query_params: QueryPoints) -> Self {
+        Self {
+            client,
+            model,
+            query_params,
+        }
+    }
+
+    async fn generate_query_vector(&self, query: &str) -> Result<Vec<f32>, VectorStoreError> {
+        let embedding = self.model.embed_document(query).await?;
+        Ok(embedding.vec.iter().map(|&x| x as f32).collect())
+    }
+
+    fn prepare_query_params(&self, query: Option<Query>, limit: usize) -> QueryPoints {
+        let mut params = self.query_params.clone();
+        params.query = query;
+        params.limit = Some(limit as u64);
+        params
+    }
+}
+
+/// Converts a PointId to its string representation.
+fn stringify_id(id: PointId) -> Result<String, VectorStoreError> {
+    match id.point_id_options {
+        Some(PointIdOptions::Num(num)) => Ok(num.to_string()),
+        Some(PointIdOptions::Uuid(uuid)) => Ok(uuid.to_string()),
+        None => Err(VectorStoreError::DatastoreError(
+            "Invalid point ID format".into(),
+        )),
+    }
+}
+
+impl<M: EmbeddingModel + std::marker::Sync + Send> VectorStoreIndex for QdrantVectorStore<M> {
+    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
+        &self,
+        query: &str,
+        n: usize,
+    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
+        let query = match self.query_params.query {
+            Some(ref q) => Some(q.clone()),
+            None => Some(Query::new_nearest(self.generate_query_vector(query).await?)),
+        };
+
+        let params = self.prepare_query_params(query, n);
+        let result = self
+            .client
+            .query(params)
+            .await
+            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
+
+        result
+            .result
+            .into_iter()
+            .map(|item| {
+                let id =
+                    stringify_id(item.id.ok_or_else(|| {
+                        VectorStoreError::DatastoreError("Missing point ID".into())
+                    })?)?;
+                let score = item.score as f64;
+                let payload = serde_json::from_value(serde_json::to_value(item.payload)?)?;
+                Ok((score, id, payload))
+            })
+            .collect()
+    }
+
+    async fn top_n_ids(
+        &self,
+        query: &str,
+        n: usize,
+    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
+        let query = match self.query_params.query {
+            Some(ref q) => Some(q.clone()),
+            None => Some(Query::new_nearest(self.generate_query_vector(query).await?)),
+        };
+
+        let params = self.prepare_query_params(query, n);
+        let points = self
+            .client
+            .query(params)
+            .await
+            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?
+            .result;
+
+        points
+            .into_iter()
+            .map(|point| {
+                let id =
+                    stringify_id(point.id.ok_or_else(|| {
+                        VectorStoreError::DatastoreError("Missing point ID".into())
+                    })?)?;
+                Ok((point.score as f64, id))
+            })
+            .collect()
+    }
+}