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    async fn add_documents(
14        &self,
15        docs: &[Document],
16        opt: &VecStoreOptions,
17    ) -> Result<Vec<String>, Box<dyn Error>>;
18
19    async fn similarity_search(
20        &self,
21        query: &str,
22        limit: usize,
23        opt: &VecStoreOptions,
24    ) -> Result<Vec<Document>, Box<dyn Error>>;
25}
26impl<VS> From<VS> for Box<dyn VectorStore>
27where
28    VS: 'static + VectorStore,
29{
30    fn from(vector_store: VS) -> Self {
31        Box::new(vector_store)
32    }
33}
34
35#[macro_export]
36macro_rules! add_documents {
37    ($obj:expr, $docs:expr) => {
38        $obj.add_documents($docs, &$crate::vectorstore::VecStoreOptions::default())
39    };
40    ($obj:expr, $docs:expr, $opt:expr) => {
41        $obj.add_documents($docs, $opt)
42    };
43}
44
45#[macro_export]
46macro_rules! similarity_search {
47    ($obj:expr, $query:expr, $limit:expr) => {
48        $obj.similarity_search(
49            $query,
50            $limit,
51            &$crate::vectorstore::VecStoreOptions::default(),
52        )
53    };
54    ($obj:expr, $query:expr, $limit:expr, $opt:expr) => {
55        $obj.similarity_search($query, $limit, $opt)
56    };
57}
58
59// Retriever is a retriever for vector stores.
60pub struct Retriever {
61    vstore: Box<dyn VectorStore>,
62    num_docs: usize,
63    options: VecStoreOptions,
64}
65impl Retriever {
66    pub fn new<V: Into<Box<dyn VectorStore>>>(vstore: V, num_docs: usize) -> Self {
67        Retriever {
68            vstore: vstore.into(),
69            num_docs,
70            options: VecStoreOptions::default(),
71        }
72    }
73
74    pub fn with_options(mut self, options: VecStoreOptions) -> Self {
75        self.options = options;
76        self
77    }
78}
79
80#[async_trait]
81impl schemas::Retriever for Retriever {
82    async fn get_relevant_documents(&self, query: &str) -> Result<Vec<Document>, Box<dyn Error>> {
83        self.vstore
84            .similarity_search(query, self.num_docs, &self.options)
85            .await
86    }
87}
langchain_rust/vectorstore/vectorstore.rs

langchain_rust/vectorstore/
vectorstore.rs
