rig_milvus/
lib.rs

1mod filter;
2
3use reqwest::StatusCode;
4use rig::{
5    Embed, OneOrMany,
6    embeddings::{Embedding, EmbeddingModel},
7    vector_store::{
8        InsertDocuments, VectorStoreError, VectorStoreIndex,
9        request::{SearchFilter, VectorSearchRequest},
10    },
11};
12use serde::{Deserialize, Serialize};
13
14use crate::filter::Filter;
15
16/// Represents a vector store implementation using Milvus - <https://milvus.io/> as the backend.
17pub struct MilvusVectorStore<M> {
18    /// Model used to generate embeddings for the vector store
19    model: M,
20    base_url: String,
21    client: reqwest::Client,
22    database_name: String,
23    collection_name: String,
24    token: Option<String>,
25}
26
27#[derive(Debug, Serialize, Deserialize)]
28pub struct CreateRecord {
29    document: String,
30    embedded_text: String,
31    embedding: Vec<f64>,
32}
33
34#[derive(Debug, Serialize, Deserialize)]
35#[serde(rename_all = "camelCase")]
36struct InsertRequest<'a> {
37    data: Vec<CreateRecord>,
38    collection_name: &'a str,
39    db_name: &'a str,
40}
41
42#[derive(Debug, Serialize, Deserialize)]
43#[serde(rename_all = "camelCase")]
44struct SearchRequest<'a> {
45    collection_name: &'a str,
46    db_name: &'a str,
47    data: Vec<f64>,
48    #[serde(skip_serializing_if = "String::is_empty")]
49    filter: String,
50    anns_field: &'a str,
51    limit: usize,
52    output_fields: Vec<&'a str>,
53}
54
55#[derive(Debug, Serialize, Deserialize)]
56#[serde(rename_all = "camelCase")]
57struct SearchResult<T> {
58    code: i64,
59    data: Vec<SearchResultData<T>>,
60}
61
62#[derive(Debug, Serialize, Deserialize)]
63#[serde(rename_all = "camelCase")]
64struct SearchResultData<T> {
65    id: i64,
66    distance: f64,
67    document: T,
68    embedded_text: String,
69}
70
71#[derive(Debug, Serialize, Deserialize)]
72#[serde(rename_all = "camelCase")]
73struct SearchResultOnlyId {
74    code: i64,
75    data: Vec<SearchResultDataOnlyId>,
76}
77
78#[derive(Debug, Serialize, Deserialize)]
79#[serde(rename_all = "camelCase")]
80struct SearchResultDataOnlyId {
81    id: i64,
82    distance: f64,
83}
84
85impl<M> MilvusVectorStore<M>
86where
87    M: EmbeddingModel,
88{
89    /// Creates a new instance of `MilvusVectorStore`.
90    ///
91    /// # Arguments
92    /// * `model` - Embedding model instance
93    /// * `base_url` - The URL of where your Milvus instance is located. Alternatively if you're using the Milvus offering provided by Zilliz, your cluster endpoint.
94    /// * `database_name` - The name of your database
95    /// * `collection_name` - The name of your collection
96    pub fn new(model: M, base_url: String, database_name: String, collection_name: String) -> Self {
97        Self {
98            model,
99            base_url,
100            client: reqwest::Client::new(),
101            database_name,
102            collection_name,
103            token: None,
104        }
105    }
106
107    /// Forms the auth token for Milvus from your username and password. Required if using a Milvus instance that requires authentication.
108    pub fn auth(mut self, username: String, password: String) -> Self {
109        let str = format!("{username}:{password}");
110        self.token = Some(str);
111
112        self
113    }
114
115    /// Creates a Milvus insertion request.
116    fn create_insert_request(&self, data: Vec<CreateRecord>) -> InsertRequest<'_> {
117        InsertRequest {
118            data,
119            collection_name: &self.collection_name,
120            db_name: &self.database_name,
121        }
122    }
123
124    /// Creates a Milvus semantic search request.
125    fn create_search_request(
126        &self,
127        data: Vec<f64>,
128        req: &VectorSearchRequest<Filter>,
129        id_only: bool,
130    ) -> SearchRequest<'_> {
131        const OUTPUT_FIELDS: [&str; 4] = ["id", "distance", "document", "embeddedText"];
132        const OUTPUT_FIELDS_ID_ONLY: [&str; 2] = ["id", "distance"];
133
134        let output_fields = if id_only {
135            OUTPUT_FIELDS_ID_ONLY.to_vec()
136        } else {
137            OUTPUT_FIELDS.to_vec()
138        };
139
140        let threshold = req
141            .threshold()
142            .map(|thresh| Filter::gte("distance".into(), thresh.into()));
143
144        let filter = match (threshold, req.filter()) {
145            (Some(thresh), Some(filter)) => thresh.and(filter.clone()).into_inner(),
146            (Some(thresh), _) => thresh.into_inner(),
147            (_, Some(filter)) => filter.clone().into_inner(),
148            _ => String::new(),
149        };
150
151        SearchRequest {
152            collection_name: &self.collection_name,
153            db_name: &self.database_name,
154            data,
155            filter,
156            anns_field: "embedding",
157            limit: req.samples() as usize,
158            output_fields,
159        }
160    }
161}
162
163impl<Model> InsertDocuments for MilvusVectorStore<Model>
164where
165    Model: EmbeddingModel + Send + Sync,
166{
167    async fn insert_documents<Doc: Serialize + Embed + Send>(
168        &self,
169        documents: Vec<(Doc, OneOrMany<Embedding>)>,
170    ) -> Result<(), VectorStoreError> {
171        let url = format!(
172            "{base_url}/v2/vectordb/entities/insert",
173            base_url = self.base_url
174        );
175
176        let data = documents
177            .into_iter()
178            .map(|(document, embeddings)| {
179                let json_document: serde_json::Value = serde_json::to_value(&document)?;
180                let json_document_as_string = serde_json::to_string(&json_document)?;
181
182                let embeddings = embeddings
183                    .into_iter()
184                    .map(|embedding| {
185                        let embedded_text = embedding.document;
186                        let embedding: Vec<f64> = embedding.vec;
187
188                        CreateRecord {
189                            document: json_document_as_string.clone(),
190                            embedded_text,
191                            embedding,
192                        }
193                    })
194                    .collect::<Vec<CreateRecord>>();
195                Ok(embeddings)
196            })
197            .collect::<Result<Vec<Vec<CreateRecord>>, VectorStoreError>>()?
198            .into_iter()
199            .flatten()
200            .collect::<Vec<CreateRecord>>();
201
202        let mut client = self.client.post(url);
203        if let Some(ref token) = self.token {
204            client = client.header("Authentication", format!("Bearer {token}"));
205        }
206
207        let insert_request = self.create_insert_request(data);
208
209        let body = serde_json::to_string(&insert_request).unwrap();
210
211        let res = client.body(body).send().await?;
212
213        if res.status() != StatusCode::OK {
214            let status = res.status();
215            let text = res.text().await?;
216
217            return Err(VectorStoreError::ExternalAPIError(status, text));
218        }
219
220        Ok(())
221    }
222}
223
224impl<M> VectorStoreIndex for MilvusVectorStore<M>
225where
226    M: EmbeddingModel,
227{
228    type Filter = Filter;
229
230    /// Search for the top `n` nearest neighbors to the given query within the Milvus vector store.
231    /// Returns a vector of tuples containing the score, ID, and payload of the nearest neighbors.
232    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
233        &self,
234        req: VectorSearchRequest<Filter>,
235    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
236        let embedding = self.model.embed_text(req.query()).await?;
237        let url = format!(
238            "{base_url}/v2/vectordb/entities/search",
239            base_url = self.base_url
240        );
241
242        let body = self.create_search_request(embedding.vec, &req, false);
243
244        let mut client = self.client.post(url);
245        if let Some(ref token) = self.token {
246            client = client.header("Authentication", format!("Bearer {token}"));
247        }
248
249        let body = serde_json::to_string(&body)?;
250
251        let res = client.body(body).send().await?;
252
253        if res.status() != StatusCode::OK {
254            let status = res.status();
255            let text = res.text().await?;
256
257            return Err(VectorStoreError::ExternalAPIError(status, text));
258        }
259
260        let json: SearchResult<T> = res.json().await?;
261
262        let res = json
263            .data
264            .into_iter()
265            .map(|x| (x.distance, x.id.to_string(), x.document))
266            .collect();
267
268        Ok(res)
269    }
270
271    /// Search for the top `n` nearest neighbors to the given query within the Milvus vector store.
272    /// Returns a vector of tuples containing the score and ID of the nearest neighbors.
273    async fn top_n_ids(
274        &self,
275        req: VectorSearchRequest<Filter>,
276    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
277        let embedding = self.model.embed_text(req.query()).await?;
278        let url = format!(
279            "{base_url}/v2/vectordb/entities/search",
280            base_url = self.base_url
281        );
282
283        let body = self.create_search_request(embedding.vec, &req, true);
284
285        let mut client = self.client.post(url);
286        if let Some(ref token) = self.token {
287            client = client.header("Authentication", format!("Bearer {token}"));
288        }
289
290        let body = serde_json::to_string(&body)?;
291
292        let res = client.body(body).send().await?;
293
294        if res.status() != StatusCode::OK {
295            let status = res.status();
296            let text = res.text().await?;
297
298            return Err(VectorStoreError::ExternalAPIError(status, text));
299        }
300
301        let json: SearchResultOnlyId = res.json().await?;
302
303        let res = json
304            .data
305            .into_iter()
306            .map(|x| (x.distance, x.id.to_string()))
307            .collect();
308
309        Ok(res)
310    }
311}
rig_milvus/lib.rs

rig_milvus/
lib.rs
