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

rig_milvus/
lib.rs
