swiftide_core/search_strategies/
similarity_single_embedding.rs

1use crate::querying;
2
3use super::{SearchFilter, DEFAULT_TOP_K};
4
5/// A simple, single vector similarity search where it takes the embedding on the current query
6/// and returns `top_k` documents.
7///
8/// Can optionally be used with a filter.
9#[derive(Debug, Clone)]
10pub struct SimilaritySingleEmbedding<FILTER: SearchFilter = ()> {
11    /// Maximum number of documents to return
12    top_k: u64,
13
14    filter: Option<FILTER>,
15}
16
17impl<FILTER: SearchFilter> querying::SearchStrategy for SimilaritySingleEmbedding<FILTER> {}
18
19impl<FILTER: SearchFilter> Default for SimilaritySingleEmbedding<FILTER> {
20    fn default() -> Self {
21        Self {
22            top_k: DEFAULT_TOP_K,
23            filter: None,
24        }
25    }
26}
27
28impl SimilaritySingleEmbedding<()> {
29    /// Set an optional filter to be used in the query
30    pub fn into_concrete_filter<FILTER: SearchFilter>(&self) -> SimilaritySingleEmbedding<FILTER> {
31        SimilaritySingleEmbedding::<FILTER> {
32            top_k: self.top_k,
33            filter: None,
34        }
35    }
36}
37
38impl<FILTER: SearchFilter> SimilaritySingleEmbedding<FILTER> {
39    pub fn from_filter(filter: FILTER) -> Self {
40        Self {
41            filter: Some(filter),
42            ..Default::default()
43        }
44    }
45
46    /// Set the maximum amount of documents to be returned
47    pub fn with_top_k(&mut self, top_k: u64) -> &mut Self {
48        self.top_k = top_k;
49
50        self
51    }
52
53    /// Returns the maximum of documents to be returned
54    pub fn top_k(&self) -> u64 {
55        self.top_k
56    }
57
58    /// Set an optional filter to be used in the query
59    pub fn with_filter<NEWFILTER: SearchFilter>(
60        self,
61        filter: NEWFILTER,
62    ) -> SimilaritySingleEmbedding<NEWFILTER> {
63        SimilaritySingleEmbedding::<NEWFILTER> {
64            top_k: self.top_k,
65            filter: Some(filter),
66        }
67    }
68
69    pub fn filter(&self) -> &Option<FILTER> {
70        &self.filter
71    }
72}
swiftide_core/search_strategies/similarity_single_embedding.rs

swiftide_core/search_strategies/
similarity_single_embedding.rs
