Small commit to add hybrid search and autoembedding

milli/src/update/settings.rs

@@ -3,7 +3,7 @@ use std::result::Result as StdResult;
 
 use charabia::{Normalize, Tokenizer, TokenizerBuilder};
 use deserr::{DeserializeError, Deserr};
-use itertools::Itertools;
+use itertools::{EitherOrBoth, Itertools};
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use time::OffsetDateTime;
 
@@ -15,6 +15,8 @@ use crate::index::{DEFAULT_MIN_WORD_LEN_ONE_TYPO, DEFAULT_MIN_WORD_LEN_TWO_TYPOS
 use crate::proximity::ProximityPrecision;
 use crate::update::index_documents::IndexDocumentsMethod;
 use crate::update::{IndexDocuments, UpdateIndexingStep};
+use crate::vector::settings::{EmbeddingSettings, PromptSettings};
+use crate::vector::EmbeddingConfig;
 use crate::{FieldsIdsMap, Index, OrderBy, Result};
 
 #[derive(Debug, Clone, PartialEq, Eq, Copy)]
@@ -73,6 +75,13 @@ impl<T> Setting<T> {
             otherwise => otherwise,
         }
     }
+
+    pub fn apply(&mut self, new: Self) {
+        if let Setting::NotSet = new {
+            return;
+        }
+        *self = new;
+    }
 }
 
 impl<T: Serialize> Serialize for Setting<T> {
@@ -129,6 +138,7 @@ pub struct Settings<'a, 't, 'i> {
     sort_facet_values_by: Setting<HashMap<String, OrderBy>>,
     pagination_max_total_hits: Setting<usize>,
     proximity_precision: Setting<ProximityPrecision>,
+    embedder_settings: Setting<BTreeMap<String, Setting<EmbeddingSettings>>>,
 }
 
 impl<'a, 't, 'i> Settings<'a, 't, 'i> {
@@ -161,6 +171,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
             sort_facet_values_by: Setting::NotSet,
             pagination_max_total_hits: Setting::NotSet,
             proximity_precision: Setting::NotSet,
+            embedder_settings: Setting::NotSet,
             indexer_config,
         }
     }
@@ -343,6 +354,14 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
         self.proximity_precision = Setting::Reset;
     }
 
+    pub fn set_embedder_settings(&mut self, value: BTreeMap<String, Setting<EmbeddingSettings>>) {
+        self.embedder_settings = Setting::Set(value);
+    }
+
+    pub fn reset_embedder_settings(&mut self) {
+        self.embedder_settings = Setting::Reset;
+    }
+
     fn reindex<FP, FA>(
         &mut self,
         progress_callback: &FP,
@@ -890,6 +909,60 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
         Ok(changed)
     }
 
+    fn update_embedding_configs(&mut self) -> Result<bool> {
+        let update = match std::mem::take(&mut self.embedder_settings) {
+            Setting::Set(configs) => {
+                let mut changed = false;
+                let old_configs = self.index.embedding_configs(self.wtxn)?;
+                let old_configs: BTreeMap<String, Setting<EmbeddingSettings>> =
+                    old_configs.into_iter().map(|(k, v)| (k, Setting::Set(v.into()))).collect();
+
+                let mut new_configs = BTreeMap::new();
+                for joined in old_configs
+                    .into_iter()
+                    .merge_join_by(configs.into_iter(), |(left, _), (right, _)| left.cmp(right))
+                {
+                    match joined {
+                        EitherOrBoth::Both((name, mut old), (_, new)) => {
+                            old.apply(new);
+                            let new = validate_prompt(&name, old)?;
+                            changed = true;
+                            new_configs.insert(name, new);
+                        }
+                        EitherOrBoth::Left((name, setting)) => {
+                            new_configs.insert(name, setting);
+                        }
+                        EitherOrBoth::Right((name, setting)) => {
+                            let setting = validate_prompt(&name, setting)?;
+                            changed = true;
+                            new_configs.insert(name, setting);
+                        }
+                    }
+                }
+                let new_configs: Vec<(String, EmbeddingConfig)> = new_configs
+                    .into_iter()
+                    .filter_map(|(name, setting)| match setting {
+                        Setting::Set(value) => Some((name, value.into())),
+                        Setting::Reset => None,
+                        Setting::NotSet => Some((name, EmbeddingSettings::default().into())),
+                    })
+                    .collect();
+                if new_configs.is_empty() {
+                    self.index.delete_embedding_configs(self.wtxn)?;
+                } else {
+                    self.index.put_embedding_configs(self.wtxn, new_configs)?;
+                }
+                changed
+            }
+            Setting::Reset => {
+                self.index.delete_embedding_configs(self.wtxn)?;
+                true
+            }
+            Setting::NotSet => false,
+        };
+        Ok(update)
+    }
+
     pub fn execute<FP, FA>(mut self, progress_callback: FP, should_abort: FA) -> Result<()>
     where
         FP: Fn(UpdateIndexingStep) + Sync,
@@ -927,6 +1000,13 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
         let searchable_updated = self.update_searchable()?;
         let exact_attributes_updated = self.update_exact_attributes()?;
         let proximity_precision = self.update_proximity_precision()?;
+        // TODO: very rough approximation of the needs for reindexing where any change will result in
+        // a full reindexing.
+        // What can be done instead:
+        // 1. Only change the distance on a distance change
+        // 2. Only change the name -> embedder mapping on a name change
+        // 3. Keep the old vectors but reattempt indexing on a prompt change: only actually changed prompt will need embedding + storage
+        let embedding_configs_updated = self.update_embedding_configs()?;
 
         if stop_words_updated
             || non_separator_tokens_updated
@@ -937,6 +1017,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
             || searchable_updated
             || exact_attributes_updated
             || proximity_precision
+            || embedding_configs_updated
         {
             self.reindex(&progress_callback, &should_abort, old_fields_ids_map)?;
         }
@@ -945,6 +1026,34 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> {
     }
 }
 
+fn validate_prompt(
+    name: &str,
+    new: Setting<EmbeddingSettings>,
+) -> Result<Setting<EmbeddingSettings>> {
+    match new {
+        Setting::Set(EmbeddingSettings {
+            embedder_options,
+            prompt:
+                Setting::Set(PromptSettings { template: Setting::Set(template), strategy, fallback }),
+        }) => {
+            // validate
+            let template = crate::prompt::Prompt::new(template, None, None)
+                .map(|prompt| crate::prompt::PromptData::from(prompt).template)
+                .map_err(|inner| UserError::InvalidPromptForEmbeddings(name.to_owned(), inner))?;
+
+            Ok(Setting::Set(EmbeddingSettings {
+                embedder_options,
+                prompt: Setting::Set(PromptSettings {
+                    template: Setting::Set(template),
+                    strategy,
+                    fallback,
+                }),
+            }))
+        }
+        new => Ok(new),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use big_s::S;
@@ -1763,6 +1872,7 @@ mod tests {
                     sort_facet_values_by,
                     pagination_max_total_hits,
                     proximity_precision,
+                    embedder_settings,
                 } = settings;
                 assert!(matches!(searchable_fields, Setting::NotSet));
                 assert!(matches!(displayed_fields, Setting::NotSet));
@@ -1785,6 +1895,7 @@ mod tests {
                 assert!(matches!(sort_facet_values_by, Setting::NotSet));
                 assert!(matches!(pagination_max_total_hits, Setting::NotSet));
                 assert!(matches!(proximity_precision, Setting::NotSet));
+                assert!(matches!(embedder_settings, Setting::NotSet));
             })
             .unwrap();
     }