1#[macro_use]
12mod document;
13
14use aws_sdk_s3vectors::{
15 Client,
16 types::{PutInputVector, VectorData},
17};
18use aws_smithy_types::Document;
19use rig_core::{
20 embeddings::EmbeddingModel,
21 vector_store::{
22 InsertDocuments, VectorStoreError, VectorStoreIndex,
23 request::{SearchFilter, VectorSearchRequest},
24 },
25};
26use serde::{Deserialize, Serialize};
27use serde_json::Value;
28use std::collections::HashMap;
29use uuid::Uuid;
30
31#[derive(Debug, Serialize, Deserialize)]
32pub struct CreateRecord {
33 document: serde_json::Value,
34 embedded_text: String,
35}
36
37#[derive(Clone, Debug)]
39pub struct S3SearchFilter(aws_smithy_types::Document);
40
41impl SearchFilter for S3SearchFilter {
42 type Value = aws_smithy_types::Document;
43
44 fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
45 let key = key.as_ref().to_owned();
46 Self(document!({ key: { "$eq": value } }))
47 }
48
49 fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
50 let key = key.as_ref().to_owned();
51 Self(document!({ key: { "$gt": value } }))
52 }
53
54 fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
55 let key = key.as_ref().to_owned();
56 Self(document!({ key: { "$lt": value } }))
57 }
58
59 fn and(self, rhs: Self) -> Self {
60 Self(document!({ "$and": [ self.0, rhs.0 ]}))
61 }
62
63 fn or(self, rhs: Self) -> Self {
64 Self(document!({ "$or": [ self.0, rhs.0 ]}))
65 }
66}
67
68impl S3SearchFilter {
69 pub fn inner(&self) -> &aws_smithy_types::Document {
70 &self.0
71 }
72
73 pub fn into_inner(self) -> aws_smithy_types::Document {
74 self.0
75 }
76
77 pub fn gte(key: String, value: <Self as SearchFilter>::Value) -> Self {
78 Self(document!({ key: { "$gte": value } }))
79 }
80
81 pub fn lte(key: String, value: <Self as SearchFilter>::Value) -> Self {
82 Self(document!({ key: { "$lte": value } }))
83 }
84
85 pub fn exists(key: String) -> Self {
86 Self(document!({ "$exists": { key: true } }))
87 }
88
89 #[allow(clippy::should_implement_trait)]
90 pub fn not(self) -> Self {
91 Self(document!({ "$not": self.0 }))
92 }
93}
94
95pub struct S3VectorsVectorStore<M> {
96 embedding_model: M,
97 client: Client,
98 bucket_name: String,
99 index_name: String,
100}
101
102impl<M> S3VectorsVectorStore<M>
103where
104 M: EmbeddingModel,
105{
106 pub fn new(
107 embedding_model: M,
108 client: aws_sdk_s3vectors::Client,
109 bucket_name: &str,
110 index_name: &str,
111 ) -> Self {
112 Self {
113 embedding_model,
114 client,
115 bucket_name: bucket_name.to_string(),
116 index_name: index_name.to_string(),
117 }
118 }
119
120 pub fn bucket_name(&self) -> &str {
121 &self.bucket_name
122 }
123
124 pub fn set_bucket_name(&mut self, bucket_name: &str) {
125 self.bucket_name = bucket_name.to_string();
126 }
127
128 pub fn index_name(&self) -> &str {
129 &self.index_name
130 }
131
132 pub fn set_index_name(&mut self, index_name: &str) {
133 self.index_name = index_name.to_string();
134 }
135
136 pub fn client(&self) -> &Client {
137 &self.client
138 }
139}
140
141impl<M> InsertDocuments for S3VectorsVectorStore<M>
142where
143 M: EmbeddingModel,
144{
145 async fn insert_documents<Doc: serde::Serialize + rig_core::Embed + Send>(
146 &self,
147 documents: Vec<(Doc, rig_core::OneOrMany<rig_core::embeddings::Embedding>)>,
148 ) -> Result<(), rig_core::vector_store::VectorStoreError> {
149 let docs: Vec<PutInputVector> = documents
150 .into_iter()
151 .map(|x| {
152 let json_value = serde_json::to_value(&x.0).map_err(VectorStoreError::JsonError)?;
153
154 x.1.into_iter()
155 .map(|y| {
156 let document = CreateRecord {
157 document: json_value.clone(),
158 embedded_text: y.document,
159 };
160 let document =
161 serde_json::to_value(&document).map_err(VectorStoreError::JsonError)?;
162 let document = json_value_to_document(&document);
163 let vec = y.vec.into_iter().map(|item| item as f32).collect();
164 PutInputVector::builder()
165 .metadata(document.clone())
166 .data(VectorData::Float32(vec))
167 .key(Uuid::new_v4())
168 .build()
169 .map_err(|x| {
170 VectorStoreError::DatastoreError(
171 format!("Couldn't build vector input: {x}").into(),
172 )
173 })
174 })
175 .collect()
176 })
177 .collect::<Result<Vec<Vec<PutInputVector>>, VectorStoreError>>()
178 .map_err(|x| {
179 VectorStoreError::DatastoreError(
180 format!("Could not build vector store data: {x}").into(),
181 )
182 })?
183 .into_iter()
184 .flatten()
185 .collect();
186
187 self.client
188 .put_vectors()
189 .vector_bucket_name(self.bucket_name())
190 .set_vectors(Some(docs))
191 .set_index_name(Some(self.index_name.clone()))
192 .send()
193 .await
194 .map_err(|x| {
195 VectorStoreError::DatastoreError(
196 format!("Error while submitting document insertion request: {x}").into(),
197 )
198 })?;
199
200 Ok(())
201 }
202}
203
204fn json_value_to_document(value: &Value) -> Document {
205 match value {
206 Value::Null => Document::Null,
207 Value::Bool(b) => Document::Bool(*b),
208 Value::Number(n) => {
209 if let Some(i) = n.as_i64() {
210 Document::Number(aws_smithy_types::Number::NegInt(i))
211 } else if let Some(u) = n.as_u64() {
212 Document::Number(aws_smithy_types::Number::PosInt(u))
213 } else if let Some(f) = n.as_f64() {
214 Document::Number(aws_smithy_types::Number::Float(f))
215 } else {
216 Document::Null }
218 }
219 Value::String(s) => Document::String(s.clone()),
220 Value::Array(arr) => Document::Array(arr.iter().map(json_value_to_document).collect()),
221 Value::Object(obj) => Document::Object(
222 obj.iter()
223 .map(|(k, v)| (k.clone(), json_value_to_document(v)))
224 .collect::<HashMap<_, _>>(),
225 ),
226 }
227}
228
229fn document_to_json_value(value: &Document) -> Value {
230 match value {
231 Document::Null => Value::Null,
232 Document::Bool(b) => Value::Bool(*b),
233 Document::Number(n) => match n {
234 aws_smithy_types::Number::Float(f) => serde_json::Number::from_f64(*f)
235 .map(Value::Number)
236 .unwrap_or_else(|| Value::String(f.to_string())),
237 aws_smithy_types::Number::NegInt(i) => {
238 serde_json::Value::Number(serde_json::Number::from(*i))
239 }
240 aws_smithy_types::Number::PosInt(u) => {
241 serde_json::Value::Number(serde_json::Number::from(*u))
242 }
243 },
244 Document::String(s) => Value::String(s.clone()),
245 Document::Array(arr) => Value::Array(arr.iter().map(document_to_json_value).collect()),
246 Document::Object(obj) => {
247 let res = obj
248 .iter()
249 .map(|(k, v)| (k.clone(), document_to_json_value(v)))
250 .collect::<serde_json::Map<String, serde_json::Value>>();
251
252 serde_json::Value::Object(res)
253 }
254 }
255}
256
257impl<M> VectorStoreIndex for S3VectorsVectorStore<M>
258where
259 M: EmbeddingModel,
260{
261 type Filter = S3SearchFilter;
262
263 async fn top_n<T: for<'a> serde::Deserialize<'a> + Send>(
264 &self,
265 req: VectorSearchRequest<S3SearchFilter>,
266 ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
267 if req.samples() > i32::MAX as u64 {
268 return Err(VectorStoreError::DatastoreError(format!("The number of samples to return with the `rig` AWS S3Vectors integration cannot be higher than {}", i32::MAX).into()));
269 }
270
271 let embedding = self
272 .embedding_model
273 .embed_text(req.query())
274 .await?
275 .vec
276 .into_iter()
277 .map(|x| x as f32)
278 .collect();
279
280 let mut query_builder = self
281 .client
282 .query_vectors()
283 .query_vector(VectorData::Float32(embedding))
284 .top_k(req.samples() as i32)
285 .return_distance(true)
286 .return_metadata(true)
287 .vector_bucket_name(self.bucket_name())
288 .index_name(self.index_name());
289
290 if let Some(filter) = req.filter() {
291 query_builder = query_builder.filter(filter.inner().clone())
292 }
293
294 let query = query_builder
295 .send()
296 .await
297 .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
298
299 let res: Vec<(f64, String, T)> = query
300 .vectors
301 .into_iter()
302 .map(|x| {
303 let distance = x.distance.ok_or_else(|| {
304 VectorStoreError::DatastoreError(Box::new(std::io::Error::other(
305 "S3Vectors response missing distance",
306 )))
307 })? as f64;
308
309 if req
310 .threshold()
311 .is_some_and(|threshold| distance < threshold)
312 {
313 return Ok(None);
314 }
315
316 let metadata_document = x.metadata.ok_or_else(|| {
317 VectorStoreError::DatastoreError(Box::new(std::io::Error::other(
318 "S3Vectors response missing metadata",
319 )))
320 })?;
321 let val = document_to_json_value(&metadata_document);
322 let metadata: T = serde_json::from_value(val)?;
323
324 Ok(Some((distance, x.key, metadata)))
325 })
326 .collect::<Result<Vec<_>, VectorStoreError>>()?
327 .into_iter()
328 .flatten()
329 .collect();
330
331 Ok(res)
332 }
333
334 async fn top_n_ids(
335 &self,
336 req: VectorSearchRequest<S3SearchFilter>,
337 ) -> Result<Vec<(f64, String)>, VectorStoreError> {
338 if req.samples() > i32::MAX as u64 {
339 return Err(VectorStoreError::DatastoreError(format!("The number of samples to return with the `rig` AWS S3Vectors integration cannot be higher than {}", i32::MAX).into()));
340 }
341
342 let embedding = self
343 .embedding_model
344 .embed_text(req.query())
345 .await?
346 .vec
347 .into_iter()
348 .map(|x| x as f32)
349 .collect();
350
351 let mut query_builder = self
352 .client
353 .query_vectors()
354 .query_vector(VectorData::Float32(embedding))
355 .top_k(req.samples() as i32)
356 .return_distance(true)
357 .vector_bucket_name(self.bucket_name())
358 .index_name(self.index_name());
359
360 if let Some(filter) = req.filter() {
361 query_builder = query_builder.filter(filter.inner().clone())
362 }
363
364 let query = query_builder
365 .send()
366 .await
367 .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
368
369 let res: Vec<(f64, String)> = query
370 .vectors
371 .into_iter()
372 .map(|x| {
373 let distance = x.distance.ok_or_else(|| {
374 VectorStoreError::DatastoreError(Box::new(std::io::Error::other(
375 "S3Vectors response missing distance",
376 )))
377 })? as f64;
378
379 if req
380 .threshold()
381 .is_some_and(|threshold| distance < threshold)
382 {
383 return Ok(None);
384 }
385
386 Ok(Some((distance, x.key)))
387 })
388 .collect::<Result<Vec<_>, VectorStoreError>>()?
389 .into_iter()
390 .flatten()
391 .collect();
392
393 Ok(res)
394 }
395}