swiftide_core/search_strategies/
similarity_single_embedding.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
use crate::querying;

use super::{SearchFilter, DEFAULT_TOP_K};

/// A simple, single vector similarity search where it takes the embedding on the current query
/// and returns `top_k` documents.
///
/// Can optionally be used with a filter.
#[derive(Debug, Clone)]
pub struct SimilaritySingleEmbedding<FILTER: SearchFilter = ()> {
    /// Maximum number of documents to return
    top_k: u64,

    filter: Option<FILTER>,
}

impl<FILTER: SearchFilter> querying::SearchStrategy for SimilaritySingleEmbedding<FILTER> {}

impl<FILTER: SearchFilter> Default for SimilaritySingleEmbedding<FILTER> {
    fn default() -> Self {
        Self {
            top_k: DEFAULT_TOP_K,
            filter: None,
        }
    }
}

impl SimilaritySingleEmbedding<()> {
    /// Set an optional filter to be used in the query
    pub fn into_concrete_filter<FILTER: SearchFilter>(&self) -> SimilaritySingleEmbedding<FILTER> {
        SimilaritySingleEmbedding::<FILTER> {
            top_k: self.top_k,
            filter: None,
        }
    }
}

impl<FILTER: SearchFilter> SimilaritySingleEmbedding<FILTER> {
    pub fn from_filter(filter: FILTER) -> Self {
        Self {
            filter: Some(filter),
            ..Default::default()
        }
    }

    /// Set the maximum amount of documents to be returned
    pub fn with_top_k(&mut self, top_k: u64) -> &mut Self {
        self.top_k = top_k;

        self
    }

    /// Returns the maximum of documents to be returned
    pub fn top_k(&self) -> u64 {
        self.top_k
    }

    /// Set an optional filter to be used in the query
    pub fn with_filter<NEWFILTER: SearchFilter>(
        self,
        filter: NEWFILTER,
    ) -> SimilaritySingleEmbedding<NEWFILTER> {
        SimilaritySingleEmbedding::<NEWFILTER> {
            top_k: self.top_k,
            filter: Some(filter),
        }
    }

    pub fn filter(&self) -> &Option<FILTER> {
        &self.filter
    }
}