Small commit to add hybrid search and autoembedding

index-scheduler/src/batch.rs

@@ -1202,6 +1202,10 @@ impl IndexScheduler {
 
                 let config = IndexDocumentsConfig { update_method: method, ..Default::default() };
 
+                let embedder_configs = index.embedding_configs(index_wtxn)?;
+                // TODO: consider Arc'ing the map too (we only need read access + we'll be cloning it multiple times, so really makes sense)
+                let embedders = self.embedders(embedder_configs)?;
+
                 let mut builder = milli::update::IndexDocuments::new(
                     index_wtxn,
                     index,
@@ -1220,6 +1224,8 @@ impl IndexScheduler {
                             let (new_builder, user_result) = builder.add_documents(reader)?;
                             builder = new_builder;
 
+                            builder = builder.with_embedders(embedders.clone());
+
                             let received_documents =
                                 if let Some(Details::DocumentAdditionOrUpdate {
                                     received_documents,
@@ -1345,6 +1351,9 @@ impl IndexScheduler {
 
                 for (task, (_, settings)) in tasks.iter_mut().zip(settings) {
                     let checked_settings = settings.clone().check();
+                    if matches!(checked_settings.embedders, milli::update::Setting::Set(_)) {
+                        self.features().check_vector("Passing `embedders` in settings")?
+                    }
                     if checked_settings.proximity_precision.set().is_some() {
                         self.features.features().check_proximity_precision()?;
                     }

index-scheduler/src/features.rs

@@ -56,12 +56,12 @@ impl RoFeatures {
         }
     }
 
-    pub fn check_vector(&self) -> Result<()> {
+    pub fn check_vector(&self, disabled_action: &'static str) -> Result<()> {
         if self.runtime.vector_store {
             Ok(())
         } else {
             Err(FeatureNotEnabledError {
-                disabled_action: "Passing `vector` as a query parameter",
+                disabled_action,
                 feature: "vector store",
                 issue_link: "https://github.com/meilisearch/product/discussions/677",
             }

index-scheduler/src/insta_snapshot.rs

@@ -41,6 +41,7 @@ pub fn snapshot_index_scheduler(scheduler: &IndexScheduler) -> String {
         planned_failures: _,
         run_loop_iteration: _,
         currently_updating_index: _,
+        embedders: _,
     } = scheduler;
 
     let rtxn = env.read_txn().unwrap();

index-scheduler/src/lib.rs

@@ -52,6 +52,7 @@ use meilisearch_types::heed::types::{SerdeBincode, SerdeJson, Str, I128};
 use meilisearch_types::heed::{self, Database, Env, PutFlags, RoTxn, RwTxn};
 use meilisearch_types::milli::documents::DocumentsBatchBuilder;
 use meilisearch_types::milli::update::IndexerConfig;
+use meilisearch_types::milli::vector::{Embedder, EmbedderOptions};
 use meilisearch_types::milli::{self, CboRoaringBitmapCodec, Index, RoaringBitmapCodec, BEU32};
 use meilisearch_types::tasks::{Kind, KindWithContent, Status, Task};
 use puffin::FrameView;
@@ -341,6 +342,8 @@ pub struct IndexScheduler {
     /// so that a handle to the index is available from other threads (search) in an optimized manner.
     currently_updating_index: Arc<RwLock<Option<(String, Index)>>>,
 
+    embedders: Arc<RwLock<HashMap<EmbedderOptions, std::sync::Arc<Embedder>>>>,
+
     // ================= test
     // The next entry is dedicated to the tests.
     /// Provide a way to set a breakpoint in multiple part of the scheduler.
@@ -386,6 +389,7 @@ impl IndexScheduler {
             auth_path: self.auth_path.clone(),
             version_file_path: self.version_file_path.clone(),
             currently_updating_index: self.currently_updating_index.clone(),
+            embedders: self.embedders.clone(),
             #[cfg(test)]
             test_breakpoint_sdr: self.test_breakpoint_sdr.clone(),
             #[cfg(test)]
@@ -484,6 +488,7 @@ impl IndexScheduler {
             auth_path: options.auth_path,
             version_file_path: options.version_file_path,
             currently_updating_index: Arc::new(RwLock::new(None)),
+            embedders: Default::default(),
 
             #[cfg(test)]
             test_breakpoint_sdr,
@@ -1333,6 +1338,42 @@ impl IndexScheduler {
         }
     }
 
+    // TODO: consider using a type alias or a struct embedder/template
+    #[allow(clippy::type_complexity)]
+    pub fn embedders(
+        &self,
+        embedding_configs: Vec<(String, milli::vector::EmbeddingConfig)>,
+    ) -> Result<HashMap<String, (Arc<milli::vector::Embedder>, Arc<milli::prompt::Prompt>)>> {
+        let res: Result<_> = embedding_configs
+            .into_iter()
+            .map(|(name, milli::vector::EmbeddingConfig { embedder_options, prompt })| {
+                let prompt =
+                    Arc::new(prompt.try_into().map_err(meilisearch_types::milli::Error::from)?);
+                // optimistically return existing embedder
+                {
+                    let embedders = self.embedders.read().unwrap();
+                    if let Some(embedder) = embedders.get(&embedder_options) {
+                        return Ok((name, (embedder.clone(), prompt)));
+                    }
+                }
+
+                // add missing embedder
+                let embedder = Arc::new(
+                    Embedder::new(embedder_options.clone())
+                        .map_err(meilisearch_types::milli::vector::Error::from)
+                        .map_err(meilisearch_types::milli::UserError::from)
+                        .map_err(meilisearch_types::milli::Error::from)?,
+                );
+                {
+                    let mut embedders = self.embedders.write().unwrap();
+                    embedders.insert(embedder_options, embedder.clone());
+                }
+                Ok((name, (embedder, prompt)))
+            })
+            .collect();
+        res
+    }
+
     /// Blocks the thread until the test handle asks to progress to/through this breakpoint.
     ///
     /// Two messages are sent through the channel for each breakpoint.