1use serde::{Deserialize, Serialize};
4
5use crate::prelude::*;
6
7#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
9pub struct SearchRequest {
10 pub query: String,
13
14 pub source: DocumentsSource,
16
17 #[serde(skip_serializing_if = "Option::is_none")]
19 pub instructions: Option<String>,
20
21 #[serde(skip_serializing_if = "Option::is_none")]
23 pub limit: Option<i32>,
24
25 #[serde(skip_serializing_if = "Option::is_none")]
27 pub ranking_metric: Option<RankingMetric>,
28
29 #[serde(skip_serializing_if = "Option::is_none")]
31 pub retrieval_mode: Option<RetrievalMode>,
32}
33
34#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
36pub struct SearchResponse {
37 pub matches: Vec<SearchMatch>,
39}
40
41#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
43pub struct SearchMatch {
44 pub file_id: String,
46
47 pub chunk_id: String,
49
50 pub chunk_content: String,
52
53 pub score: f32,
55
56 pub collection_ids: Vec<String>,
58
59 #[serde(default)]
61 pub fields: HashMap<String, String>,
62}
63
64#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
66pub struct DocumentsSource {
67 pub collection_ids: Vec<String>,
69}
70
71#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
73pub struct SearchParameters {
74 #[serde(skip_serializing_if = "Option::is_none")]
76 pub from_date: Option<String>,
77
78 #[serde(skip_serializing_if = "Option::is_none")]
80 pub to_date: Option<String>,
81
82 #[serde(skip_serializing_if = "Option::is_none")]
84 pub max_search_results: Option<i32>,
85
86 #[serde(skip_serializing_if = "Option::is_none")]
88 pub mode: Option<String>,
89
90 #[serde(skip_serializing_if = "Option::is_none")]
92 pub return_citations: Option<bool>,
93
94 #[serde(skip_serializing_if = "Option::is_none")]
96 pub sources: Option<Vec<SearchSource>>,
97}
98
99#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
101#[serde(tag = "type", rename_all = "lowercase")]
102pub enum SearchSource {
103 X {
105 #[serde(skip_serializing_if = "Option::is_none")]
107 excluded_x_handles: Option<Vec<String>>,
108
109 #[serde(skip_serializing_if = "Option::is_none")]
111 included_x_handles: Option<Vec<String>>,
112
113 #[serde(skip_serializing_if = "Option::is_none")]
115 x_handles: Option<Vec<String>>,
116
117 #[serde(skip_serializing_if = "Option::is_none")]
119 post_favorite_count: Option<i32>,
120
121 #[serde(skip_serializing_if = "Option::is_none")]
123 post_view_count: Option<i32>,
124 },
125 Web {
127 #[serde(skip_serializing_if = "Option::is_none")]
129 allowed_websites: Option<Vec<String>>,
130
131 #[serde(skip_serializing_if = "Option::is_none")]
133 excluded_websites: Option<Vec<String>>,
134
135 #[serde(skip_serializing_if = "Option::is_none")]
137 country: Option<String>,
138
139 #[serde(skip_serializing_if = "Option::is_none")]
141 safe_search: Option<bool>,
142 },
143 News {
145 #[serde(skip_serializing_if = "Option::is_none")]
147 country: Option<String>,
148
149 #[serde(skip_serializing_if = "Option::is_none")]
151 excluded_websites: Option<Vec<String>>,
152 },
153 Rss {
155 #[serde(skip_serializing_if = "Option::is_none")]
157 urls: Option<Vec<String>>,
158 },
159}
160
161impl Default for SearchSource {
162 fn default() -> Self {
163 SearchSource::Web {
164 allowed_websites: None,
165 excluded_websites: None,
166 country: None,
167 safe_search: None,
168 }
169 }
170}
171
172#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
174#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
175pub enum RankingMetric {
176 #[default]
178 RankingMetricUnknown,
179 RankingMetricL2Distance,
181 RankingMetricCosineSimilarity,
183}
184
185#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
187#[serde(tag = "type", rename_all = "lowercase")]
188pub enum RetrievalMode {
189 Hybrid {
191 #[serde(skip_serializing_if = "Option::is_none")]
193 reranker: Option<HybridReranker>,
194
195 #[serde(skip_serializing_if = "Option::is_none")]
197 search_multiplier: Option<i32>,
198 },
199 Semantic {
201 #[serde(skip_serializing_if = "Option::is_none")]
203 reranker: Option<RerankerModel>,
204 },
205 Keyword {
207 #[serde(skip_serializing_if = "Option::is_none")]
209 reranker: Option<RerankerModel>,
210 },
211}
212
213impl Default for RetrievalMode {
214 fn default() -> Self {
215 RetrievalMode::Hybrid {
216 reranker: None,
217 search_multiplier: None,
218 }
219 }
220}
221
222#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
224#[serde(tag = "type", rename_all = "snake_case")]
225pub enum HybridReranker {
226 RerankerModel {
228 #[serde(skip_serializing_if = "Option::is_none")]
230 instructions: Option<String>,
231
232 #[serde(skip_serializing_if = "Option::is_none")]
234 model: Option<String>,
235 },
236 Rrf {
238 #[serde(skip_serializing_if = "Option::is_none")]
240 embedding_weight: Option<f32>,
241
242 #[serde(skip_serializing_if = "Option::is_none")]
244 text_weight: Option<f32>,
245
246 #[serde(skip_serializing_if = "Option::is_none")]
248 k: Option<i32>,
249 },
250}
251
252impl Default for HybridReranker {
253 fn default() -> Self {
254 HybridReranker::Rrf {
255 embedding_weight: None,
256 text_weight: None,
257 k: None,
258 }
259 }
260}
261
262#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
264pub struct RerankerModel {
265 #[serde(skip_serializing_if = "Option::is_none")]
267 pub instructions: Option<String>,
268
269 #[serde(skip_serializing_if = "Option::is_none")]
271 pub model: Option<String>,
272}
273
274#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
276pub struct ReciprocalRankFusion {
277 #[serde(skip_serializing_if = "Option::is_none")]
279 pub embedding_weight: Option<f32>,
280
281 #[serde(skip_serializing_if = "Option::is_none")]
283 pub text_weight: Option<f32>,
284
285 #[serde(skip_serializing_if = "Option::is_none")]
287 pub k: Option<i32>,
288}