rig_milvus/
lib.rs

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

rig_milvus/
lib.rs
