langchain_rust/vectorstore/
vectorstore.rs

1use std::error::Error;
2
3use async_trait::async_trait;
4
5use crate::schemas::{self, Document};
6
7use super::VecStoreOptions;
8
9// VectorStore is the trait for saving and querying documents in the
10// form of vector embeddings.
11#[async_trait]
12pub trait VectorStore: Send + Sync {
13    type Options;
14
15    async fn add_documents(
16        &self,
17        docs: &[Document],
18        opt: &Self::Options,
19    ) -> Result<Vec<String>, Box<dyn Error>>;
20
21    async fn similarity_search(
22        &self,
23        query: &str,
24        limit: usize,
25        opt: &Self::Options,
26    ) -> Result<Vec<Document>, Box<dyn Error>>;
27}
28
29impl<VS, F> From<VS> for Box<dyn VectorStore<Options = F>>
30where
31    VS: 'static + VectorStore<Options = F>,
32{
33    fn from(vector_store: VS) -> Self {
34        Box::new(vector_store)
35    }
36}
37
38#[macro_export]
39macro_rules! add_documents {
40    ($obj:expr, $docs:expr) => {
41        $obj.add_documents($docs, &$crate::vectorstore::VecStoreOptions::default())
42    };
43    ($obj:expr, $docs:expr, $opt:expr) => {
44        $obj.add_documents($docs, $opt)
45    };
46}
47
48#[macro_export]
49macro_rules! similarity_search {
50    ($obj:expr, $query:expr, $limit:expr) => {
51        $obj.similarity_search(
52            $query,
53            $limit,
54            &$crate::vectorstore::VecStoreOptions::default(),
55        )
56    };
57    ($obj:expr, $query:expr, $limit:expr, $opt:expr) => {
58        $obj.similarity_search($query, $limit, $opt)
59    };
60}
61
62// Retriever is a retriever for vector stores.
63pub struct Retriever<F> {
64    vstore: Box<dyn VectorStore<Options = VecStoreOptions<F>>>,
65    num_docs: usize,
66    options: VecStoreOptions<F>,
67}
68impl<F> Retriever<F> {
69    pub fn new<V: Into<Box<dyn VectorStore<Options = VecStoreOptions<F>>>>>(
70        vstore: V,
71        num_docs: usize,
72    ) -> Self {
73        Retriever {
74            vstore: vstore.into(),
75            num_docs,
76            options: VecStoreOptions::<F>::new(),
77        }
78    }
79
80    pub fn with_options(mut self, options: VecStoreOptions<F>) -> Self {
81        self.options = options;
82        self
83    }
84}
85
86#[async_trait]
87impl<O: Sync + Send> schemas::Retriever for Retriever<O> {
88    async fn get_relevant_documents(&self, query: &str) -> Result<Vec<Document>, Box<dyn Error>> {
89        self.vstore
90            .similarity_search(query, self.num_docs, &self.options)
91            .await
92    }
93}
langchain_rust/vectorstore/vectorstore.rs

langchain_rust/vectorstore/
vectorstore.rs
