Skip to main content

seshat/
config.rs

1// Copyright 2019 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use uuid::Uuid;
16#[cfg(feature = "encryption")]
17use zeroize::Zeroizing;
18
19use crate::events::{EventType, RoomId};
20
21const DEFAULT_LOAD_LIMIT: usize = 20;
22
23#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
24#[serde(default)]
25/// Search configuration
26/// A search configuration allows users to limit the search to a specific room
27/// or limit the search to specific event types.
28/// The search result can be configured in various ways as well.
29pub struct SearchConfig {
30    pub(crate) limit: usize,
31    pub(crate) before_limit: usize,
32    pub(crate) after_limit: usize,
33    pub(crate) order_by_recency: bool,
34    pub(crate) room_id: Option<RoomId>,
35    pub(crate) keys: Vec<EventType>,
36    pub(crate) next_batch: Option<Uuid>,
37}
38
39impl SearchConfig {
40    /// Create a new default search configuration.
41    pub fn new() -> Self {
42        Default::default()
43    }
44
45    /// Limit the search to a specific room.
46    /// The default is to search all rooms.
47    /// # Arguments
48    ///
49    /// * `room_id` - The unique id of the room.
50    pub fn for_room(&mut self, room_id: &str) -> &mut Self {
51        self.room_id = Some(room_id.to_owned());
52        self
53    }
54
55    /// Limit the number of events that will be returned in the search result.
56    /// The default for the limit is 10.
57    /// # Arguments
58    ///
59    /// * `limit` - The max number of events to return in the search result.
60    pub fn limit(&mut self, limit: usize) -> &mut Self {
61        self.limit = limit;
62        self
63    }
64
65    /// Limit the number of events that happened before our matching event in
66    /// the search result.
67    /// The default for the limit is 0.
68    /// # Arguments
69    ///
70    /// * `limit` - The max number of contextual events to return in the search
71    ///   result.
72    pub fn before_limit(&mut self, limit: usize) -> &mut Self {
73        self.before_limit = limit;
74        self
75    }
76
77    /// Limit the number of events that happened after our matching event in
78    /// the search result.
79    /// The default for the limit is 0.
80    /// # Arguments
81    ///
82    /// * `limit` - The max number of contextual events to return in the search
83    ///   result.
84    pub fn after_limit(&mut self, limit: usize) -> &mut Self {
85        self.after_limit = limit;
86        self
87    }
88
89    /// Should the matching events be ordered by recency. The default is to
90    /// order them by the search score.
91    /// # Arguments
92    ///
93    /// * `order_by_recency` - Flag to determine if we should order by recency.
94    ///   result.
95    pub fn order_by_recency(&mut self, order_by_recency: bool) -> &mut Self {
96        self.order_by_recency = order_by_recency;
97        self
98    }
99
100    /// Set the event types that should be used as search keys.
101    ///
102    /// This limits which events will be searched for. This method can be called
103    /// multiple times to add multiple event types. The default is to search all
104    /// event types.
105    ///
106    /// # Arguments
107    ///
108    /// * `key` - The event type that should be included in the search.
109    pub fn with_key(&mut self, key: EventType) -> &mut Self {
110        self.keys.push(key);
111        self.keys.sort();
112        self.keys.dedup();
113
114        self
115    }
116
117    /// The point to return events from. If given, this should be a next_batch
118    ///   result from a previous search.
119    pub fn next_batch(&mut self, token: Uuid) -> &mut Self {
120        self.next_batch = Some(token);
121        self
122    }
123}
124
125impl Default for SearchConfig {
126    fn default() -> Self {
127        SearchConfig {
128            limit: 10,
129            before_limit: 0,
130            after_limit: 0,
131            order_by_recency: false,
132            room_id: None,
133            keys: Vec::new(),
134            next_batch: None,
135        }
136    }
137}
138
139#[derive(Debug, PartialEq, Clone)]
140#[allow(missing_docs)]
141pub enum Language {
142    Arabic,
143    Danish,
144    Dutch,
145    English,
146    Finnish,
147    French,
148    German,
149    Greek,
150    Hungarian,
151    Italian,
152    Portuguese,
153    Romanian,
154    Russian,
155    Spanish,
156    Swedish,
157    Tamil,
158    Turkish,
159    Unknown,
160}
161
162impl Language {
163    pub(crate) fn as_tokenizer_name(&self) -> String {
164        match self {
165            Language::Unknown => "default".to_owned(),
166            lang => format!("seshat_{:?}", lang),
167        }
168    }
169
170    pub(crate) fn as_tantivy(&self) -> tantivy::tokenizer::Language {
171        match self {
172            Language::Arabic => tantivy::tokenizer::Language::Arabic,
173            Language::Danish => tantivy::tokenizer::Language::Danish,
174            Language::Dutch => tantivy::tokenizer::Language::Dutch,
175            Language::English => tantivy::tokenizer::Language::English,
176            Language::Finnish => tantivy::tokenizer::Language::Finnish,
177            Language::French => tantivy::tokenizer::Language::French,
178            Language::German => tantivy::tokenizer::Language::German,
179            Language::Greek => tantivy::tokenizer::Language::Greek,
180            Language::Hungarian => tantivy::tokenizer::Language::Hungarian,
181            Language::Italian => tantivy::tokenizer::Language::Italian,
182            Language::Portuguese => tantivy::tokenizer::Language::Portuguese,
183            Language::Romanian => tantivy::tokenizer::Language::Romanian,
184            Language::Russian => tantivy::tokenizer::Language::Russian,
185            Language::Spanish => tantivy::tokenizer::Language::Spanish,
186            Language::Swedish => tantivy::tokenizer::Language::Swedish,
187            Language::Tamil => tantivy::tokenizer::Language::Tamil,
188            Language::Turkish => tantivy::tokenizer::Language::Turkish,
189            _ => panic!("Unsupported language by tantivy"),
190        }
191    }
192}
193
194impl From<&str> for Language {
195    fn from(string: &str) -> Self {
196        match string.to_lowercase().as_ref() {
197            "arabic" => Language::Arabic,
198            "danish" => Language::Danish,
199            "dutch" => Language::Dutch,
200            "english" => Language::English,
201            "finnish" => Language::Finnish,
202            "french" => Language::French,
203            "german" => Language::German,
204            "greek" => Language::Greek,
205            "hungarian" => Language::Hungarian,
206            "italian" => Language::Italian,
207            "portuguese" => Language::Portuguese,
208            "romanian" => Language::Romanian,
209            "russian" => Language::Russian,
210            "spanish" => Language::Spanish,
211            "swedish" => Language::Swedish,
212            "tamil" => Language::Tamil,
213            "turkish" => Language::Turkish,
214            _ => Language::Unknown,
215        }
216    }
217}
218
219/// Tokenizer mode for the search index.
220///
221/// This determines how text is tokenized for indexing and searching.
222#[derive(Debug, PartialEq, Clone, Default)]
223pub enum TokenizerMode {
224    /// Language-based tokenizer using stemming (default).
225    ///
226    /// Uses SimpleTokenizer with language-specific stemming.
227    /// Best for languages with word boundaries (e.g., English, German).
228    #[default]
229    LanguageBased,
230    /// N-gram tokenizer for multi-language support.
231    ///
232    /// Splits text into character n-grams of specified sizes.
233    /// Best for languages without clear word boundaries (e.g., Japanese, Chinese).
234    Ngram {
235        /// Minimum n-gram size (default: 2)
236        min_gram: usize,
237        /// Maximum n-gram size (default: 4)
238        max_gram: usize,
239    },
240}
241
242impl TokenizerMode {
243    /// Returns the tokenizer name for this mode.
244    ///
245    /// For N-gram mode, the name includes min/max gram sizes to ensure
246    /// schema mismatch detection when these values change.
247    pub(crate) fn as_tokenizer_name(&self, language: &Language) -> String {
248        match self {
249            TokenizerMode::Ngram { min_gram, max_gram } => {
250                format!("seshat_ngram_{}_{}", min_gram, max_gram)
251            }
252            TokenizerMode::LanguageBased => language.as_tokenizer_name(),
253        }
254    }
255}
256
257#[derive(Debug, PartialEq, Clone)]
258/// Configuration for the seshat database.
259pub struct Config {
260    pub(crate) language: Language,
261    pub(crate) tokenizer_mode: TokenizerMode,
262    #[cfg(feature = "encryption")]
263    pub(crate) passphrase: Option<Zeroizing<String>>,
264}
265
266impl Config {
267    /// Create a new default Seshat database configuration.
268    pub fn new() -> Self {
269        Default::default()
270    }
271
272    /// Set the indexing language.
273    ///
274    /// This is used when `TokenizerMode::LanguageBased` is active.
275    ///
276    /// # Arguments
277    ///
278    /// * `language` - The language that will be used to index messages.
279    pub fn set_language(mut self, language: &Language) -> Self {
280        self.language = language.clone();
281        self
282    }
283
284    /// Set the tokenizer mode to N-gram.
285    ///
286    /// N-gram tokenizer splits text into character n-grams, which is useful
287    /// for languages without clear word boundaries (e.g., Japanese, Chinese).
288    ///
289    /// Note: Changing the tokenizer mode on an existing database will cause
290    /// a schema mismatch error. Delete the database and recreate it to use
291    /// a different tokenizer mode.
292    ///
293    /// # Arguments
294    ///
295    /// * `min_gram` - Minimum n-gram size (typically 2)
296    /// * `max_gram` - Maximum n-gram size (typically 4)
297    pub fn use_ngram_tokenizer(mut self, min_gram: usize, max_gram: usize) -> Self {
298        self.tokenizer_mode = TokenizerMode::Ngram { min_gram, max_gram };
299        self
300    }
301
302    /// Set the tokenizer mode to language-based (default).
303    ///
304    /// This uses SimpleTokenizer with language-specific stemming.
305    pub fn use_language_based_tokenizer(mut self) -> Self {
306        self.tokenizer_mode = TokenizerMode::LanguageBased;
307        self
308    }
309
310    /// Set the passphrase of the database.
311    /// # Arguments
312    ///
313    /// * `passphrase` - The passphrase of the database.
314    #[cfg(feature = "encryption")]
315    pub fn set_passphrase<P: Into<String>>(mut self, passphrase: P) -> Self {
316        self.passphrase = Some(Zeroizing::new(passphrase.into()));
317        self
318    }
319}
320
321impl Default for Config {
322    fn default() -> Config {
323        Config {
324            language: Language::Unknown,
325            tokenizer_mode: TokenizerMode::default(),
326            #[cfg(feature = "encryption")]
327            passphrase: None,
328        }
329    }
330}
331
332#[derive(Debug, Default, Clone, Serialize, Deserialize)]
333#[allow(missing_docs)]
334pub enum LoadDirection {
335    #[default]
336    #[serde(rename = "b", alias = "backwards", alias = "backward")]
337    Backwards,
338    #[serde(rename = "f", alias = "forwards", alias = "forward")]
339    Forwards,
340}
341
342/// Configuration for the event loading methods.
343///
344/// A load configuration allows users to limit the number of events that will be
345/// loaded or to continue loading events from a specific point in the room
346/// history.
347#[derive(Debug, Clone, Serialize, Deserialize)]
348#[serde(rename_all = "camelCase")]
349pub struct LoadConfig {
350    pub(crate) room_id: String,
351    pub(crate) limit: usize,
352    pub(crate) from_event: Option<String>,
353    #[serde(default)]
354    pub(crate) direction: LoadDirection,
355}
356
357impl LoadConfig {
358    /// Create a new LoadConfig
359    ///
360    /// The config will be created with a default limit of 20 events.
361    ///
362    /// # Arguments
363    ///
364    /// * `room_id` - The room from which to load the events.
365    pub fn new<R: Into<String>>(room_id: R) -> Self {
366        LoadConfig {
367            room_id: room_id.into(),
368            limit: DEFAULT_LOAD_LIMIT,
369            from_event: None,
370            direction: LoadDirection::default(),
371        }
372    }
373
374    /// Set the maximum amount of events that we want to load.
375    /// # Arguments
376    ///
377    /// * `limit` - The new limit that should be set.
378    pub fn limit(mut self, limit: usize) -> Self {
379        self.limit = limit;
380        self
381    }
382
383    /// Set the event from which we should continue loading events.
384    ///
385    /// # Arguments
386    ///
387    /// * `event_id` - An event id of a previous event returned by this
388    ///   method. If set events that are older than the event with the given
389    ///   event ID will be returned.
390    #[allow(clippy::wrong_self_convention)]
391    pub fn from_event<E: Into<String>>(mut self, event_id: E) -> Self {
392        self.from_event = Some(event_id.into());
393        self
394    }
395
396    /// Set the direction that we are going to continue loading events from.
397    ///
398    /// This is only used if we are continuing loading events, that is if
399    /// `from_event()` is also set.
400    ///
401    /// # Arguments
402    ///
403    /// * `direction` - The direction that should be used.
404    pub fn direction(mut self, direction: LoadDirection) -> Self {
405        self.direction = direction;
406        self
407    }
408}