Skip to main content

qdrant_client/qdrant_client/builders/
query.rs

1use crate::qdrant::{
2    ContextInput, ContextInputBuilder, ContextInputPairBuilder, DiscoverInput,
3    DiscoverInputBuilder, Formula, Mmr, NearestInputWithMmr, OrderBy, OrderByBuilder,
4    PrefetchQuery, PrefetchQueryBuilder, Query, QueryPointGroupsBuilder, QueryPointsBuilder,
5    RecommendInput, RecommendInputBuilder, RelevanceFeedbackInput, Rrf, VectorInput,
6};
7
8impl QueryPointsBuilder {
9    pub fn add_prefetch(mut self, prefetch_query: impl Into<PrefetchQuery>) -> Self {
10        self.prefetch
11            .get_or_insert_with(Vec::new)
12            .push(prefetch_query.into());
13        self
14    }
15}
16
17impl QueryPointGroupsBuilder {
18    pub fn add_prefetch(mut self, prefetch_query: impl Into<PrefetchQuery>) -> Self {
19        self.prefetch
20            .get_or_insert_with(Vec::new)
21            .push(prefetch_query.into());
22        self
23    }
24}
25
26impl PrefetchQueryBuilder {
27    pub fn add_prefetch(mut self, prefetch_query: impl Into<PrefetchQuery>) -> Self {
28        match self.prefetch {
29            Some(ref mut prefetch) => prefetch.push(prefetch_query.into()),
30            None => self.prefetch = Some(vec![prefetch_query.into()]),
31        }
32        self
33    }
34}
35
36impl Query {
37    pub fn new_nearest(value: impl Into<VectorInput>) -> Self {
38        Self {
39            variant: Some(crate::qdrant::query::Variant::Nearest(value.into())),
40        }
41    }
42
43    pub fn new_nearest_with_mmr(vector: impl Into<VectorInput>, mmr: impl Into<Mmr>) -> Self {
44        Self {
45            variant: Some(crate::qdrant::query::Variant::NearestWithMmr(
46                NearestInputWithMmr {
47                    nearest: Some(vector.into()),
48                    mmr: Some(mmr.into()),
49                },
50            )),
51        }
52    }
53
54    pub fn new_recommend(value: impl Into<RecommendInput>) -> Self {
55        Self {
56            variant: Some(crate::qdrant::query::Variant::Recommend(value.into())),
57        }
58    }
59
60    pub fn new_discover(value: impl Into<DiscoverInput>) -> Self {
61        Self {
62            variant: Some(crate::qdrant::query::Variant::Discover(value.into())),
63        }
64    }
65
66    pub fn new_context(value: impl Into<crate::qdrant::ContextInput>) -> Self {
67        Self {
68            variant: Some(crate::qdrant::query::Variant::Context(value.into())),
69        }
70    }
71
72    pub fn new_order_by(value: impl Into<OrderBy>) -> Self {
73        Self {
74            variant: Some(crate::qdrant::query::Variant::OrderBy(value.into())),
75        }
76    }
77
78    pub fn new_fusion(value: crate::qdrant::Fusion) -> Self {
79        Self {
80            variant: Some(crate::qdrant::query::Variant::Fusion(value.into())),
81        }
82    }
83
84    pub fn new_rrf(rrf: impl Into<Rrf>) -> Self {
85        Self {
86            variant: Some(crate::qdrant::query::Variant::Rrf(rrf.into())),
87        }
88    }
89
90    pub fn new_formula(formula: impl Into<Formula>) -> Self {
91        Self {
92            variant: Some(crate::qdrant::query::Variant::Formula(formula.into())),
93        }
94    }
95
96    pub fn new_sample(value: crate::qdrant::Sample) -> Self {
97        Self {
98            variant: Some(crate::qdrant::query::Variant::Sample(value.into())),
99        }
100    }
101
102    pub fn new_relevance_feedback(value: impl Into<RelevanceFeedbackInput>) -> Self {
103        Self {
104            variant: Some(crate::qdrant::query::Variant::RelevanceFeedback(
105                value.into(),
106            )),
107        }
108    }
109}
110
111impl RecommendInputBuilder {
112    pub fn add_positive(mut self, value: impl Into<VectorInput>) -> Self {
113        match self.positive {
114            Some(ref mut positive) => positive.push(value.into()),
115            None => self.positive = Some(vec![value.into()]),
116        }
117        self
118    }
119
120    pub fn add_negative(mut self, value: impl Into<VectorInput>) -> Self {
121        match self.negative {
122            Some(ref mut negative) => negative.push(value.into()),
123            None => self.negative = Some(vec![value.into()]),
124        }
125        self
126    }
127}
128
129impl DiscoverInputBuilder {
130    pub fn new(target: impl Into<VectorInput>, context: impl Into<ContextInput>) -> Self {
131        let builder = Self::empty();
132        builder.target(target).context(context)
133    }
134}
135
136impl ContextInputPairBuilder {
137    pub fn new(positive: impl Into<VectorInput>, negative: impl Into<VectorInput>) -> Self {
138        ContextInputPairBuilder::empty()
139            .positive(positive)
140            .negative(negative)
141    }
142}
143
144impl ContextInputBuilder {
145    pub fn add_pair(
146        mut self,
147        positive: impl Into<VectorInput>,
148        negative: impl Into<VectorInput>,
149    ) -> Self {
150        match self.pairs {
151            Some(ref mut pairs) => {
152                pairs.push(ContextInputPairBuilder::new(positive, negative).build())
153            }
154            None => {
155                self.pairs = Some(vec![
156                    ContextInputPairBuilder::new(positive, negative).build()
157                ])
158            }
159        }
160        self
161    }
162}
163
164impl OrderByBuilder {
165    pub fn new(key: impl Into<String>) -> Self {
166        let builder = Self::empty();
167        builder.key(key.into())
168    }
169}