rig_qdrant/
lib.rs

1//! Qdrant vector store integration for Rig.
2//!
3//! This crate provides [`QdrantVectorStore`], a Rig vector store index backed
4//! by Qdrant collections. It supports dense vector search and Qdrant filter
5//! expressions through [`QdrantFilter`].
6//!
7//! The root `rig` facade re-exports this crate as `rig::qdrant` when the
8//! `qdrant` feature is enabled.
9
10mod filter;
11
12pub use filter::QdrantFilter;
13use qdrant_client::{
14    Payload, Qdrant,
15    qdrant::{
16        Filter, PointId, PointStruct, Query, QueryPoints, UpsertPointsBuilder,
17        point_id::PointIdOptions,
18    },
19};
20use rig_core::{
21    Embed,
22    embeddings::{Embedding, EmbeddingModel},
23    vector_store::{
24        InsertDocuments, VectorStoreError, VectorStoreIndex, request::VectorSearchRequest,
25    },
26};
27use serde::{Deserialize, Serialize};
28use uuid::Uuid;
29
30/// Represents a vector store implementation using Qdrant - <https://qdrant.tech/> as the backend.
31pub struct QdrantVectorStore<M: EmbeddingModel> {
32    /// Model used to generate embeddings for the vector store
33    model: M,
34    /// Client instance for Qdrant server communication
35    client: Qdrant,
36    /// Default search parameters
37    query_params: QueryPoints,
38}
39
40impl<M> QdrantVectorStore<M>
41where
42    M: EmbeddingModel,
43{
44    /// Creates a new instance of `QdrantVectorStore`.
45    ///
46    /// # Arguments
47    /// * `client` - Qdrant client instance
48    /// * `model` - Embedding model instance
49    /// * `query_params` - Search parameters for vector queries
50    ///   Reference: <https://api.qdrant.tech/v-1-12-x/api-reference/search/query-points>
51    pub fn new(client: Qdrant, model: M, query_params: QueryPoints) -> Self {
52        Self {
53            client,
54            model,
55            query_params,
56        }
57    }
58
59    pub fn client(&self) -> &Qdrant {
60        &self.client
61    }
62
63    /// Embed query based on `QdrantVectorStore` model and modify the vector in the required format.
64    async fn generate_query_vector(&self, query: &str) -> Result<Vec<f32>, VectorStoreError> {
65        let embedding = self.model.embed_text(query).await?;
66        Ok(embedding.vec.iter().map(|&x| x as f32).collect())
67    }
68
69    /// Fill in query parameters with the given query and limit.
70    fn prepare_query_params(
71        &self,
72        query: Option<Query>,
73        limit: usize,
74        threshold: Option<f64>,
75        filter: Option<Filter>,
76    ) -> QueryPoints {
77        let mut params = self.query_params.clone();
78        params.query = query;
79        params.limit = Some(limit as u64);
80        params.score_threshold = threshold.map(|x| x as f32);
81        params.filter = filter;
82        params
83    }
84
85    /// Embeds the query (unless overridden by `query_params`), applies the
86    /// request filter, and runs the Qdrant query, returning the scored points.
87    async fn run_query(
88        &self,
89        req: &VectorSearchRequest<QdrantFilter>,
90    ) -> Result<Vec<qdrant_client::qdrant::ScoredPoint>, VectorStoreError> {
91        let query = match self.query_params.query {
92            Some(ref q) => Some(q.clone()),
93            None => Some(Query::new_nearest(
94                self.generate_query_vector(req.query()).await?,
95            )),
96        };
97
98        let filter = req
99            .filter()
100            .as_ref()
101            .cloned()
102            .map(QdrantFilter::interpret)
103            .transpose()?
104            .flatten();
105
106        let params =
107            self.prepare_query_params(query, req.samples() as usize, req.threshold(), filter);
108
109        Ok(self
110            .client
111            .query(params)
112            .await
113            .map_err(VectorStoreError::datastore)?
114            .result)
115    }
116}
117
118impl<Model> InsertDocuments for QdrantVectorStore<Model>
119where
120    Model: EmbeddingModel + Send + Sync,
121{
122    async fn insert_documents<Doc: Serialize + Embed + Send>(
123        &self,
124        documents: Vec<(Doc, Vec<Embedding>)>,
125    ) -> Result<(), VectorStoreError> {
126        let collection_name = self.query_params.collection_name.clone();
127
128        for (document, embeddings) in documents {
129            let json_document = serde_json::to_value(&document)?;
130            let doc_as_payload =
131                Payload::try_from(json_document).map_err(VectorStoreError::datastore)?;
132
133            let embeddings_as_point_structs = embeddings
134                .into_iter()
135                .map(|embedding| {
136                    let embedding_as_f32: Vec<f32> =
137                        embedding.vec.into_iter().map(|x| x as f32).collect();
138                    PointStruct::new(
139                        Uuid::new_v4().to_string(),
140                        embedding_as_f32,
141                        doc_as_payload.clone(),
142                    )
143                })
144                .collect::<Vec<PointStruct>>();
145
146            let request =
147                UpsertPointsBuilder::new(&collection_name, embeddings_as_point_structs).wait(true);
148            self.client.upsert_points(request).await.map_err(|err| {
149                VectorStoreError::DatastoreError(format!("Error while upserting: {err}").into())
150            })?;
151        }
152
153        Ok(())
154    }
155}
156
157/// Converts a `PointId` to its string representation.
158fn stringify_id(id: PointId) -> Result<String, VectorStoreError> {
159    match id.point_id_options {
160        Some(PointIdOptions::Num(num)) => Ok(num.to_string()),
161        Some(PointIdOptions::Uuid(uuid)) => Ok(uuid.to_string()),
162        None => Err(VectorStoreError::DatastoreError(
163            "Invalid point ID format".into(),
164        )),
165    }
166}
167
168impl<M> VectorStoreIndex for QdrantVectorStore<M>
169where
170    M: EmbeddingModel + std::marker::Sync + Send,
171{
172    type Filter = QdrantFilter;
173
174    /// Search for the top `n` nearest neighbors to the given query within the Qdrant vector store.
175    /// Returns a vector of tuples containing the score, ID, and payload of the nearest neighbors.
176    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
177        &self,
178        req: VectorSearchRequest<Self::Filter>,
179    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
180        self.run_query(&req)
181            .await?
182            .into_iter()
183            .map(|item| {
184                let id =
185                    stringify_id(item.id.ok_or_else(|| {
186                        VectorStoreError::DatastoreError("Missing point ID".into())
187                    })?)?;
188                let score = item.score as f64;
189                let payload = serde_json::from_value(serde_json::to_value(item.payload)?)?;
190                Ok((score, id, payload))
191            })
192            .collect()
193    }
194
195    /// Search for the top `n` nearest neighbors to the given query within the Qdrant vector store.
196    /// Returns a vector of tuples containing the score and ID of the nearest neighbors.
197    async fn top_n_ids(
198        &self,
199        req: VectorSearchRequest<Self::Filter>,
200    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
201        self.run_query(&req)
202            .await?
203            .into_iter()
204            .map(|point| {
205                let id =
206                    stringify_id(point.id.ok_or_else(|| {
207                        VectorStoreError::DatastoreError("Missing point ID".into())
208                    })?)?;
209                Ok((point.score as f64, id))
210            })
211            .collect()
212    }
213}
rig_qdrant/lib.rs

rig_qdrant/
lib.rs
