From a051ab3d9ae8ad7bf4262cbf608eb04383a6441d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Thu, 3 Jul 2025 12:04:40 +0200 Subject: [PATCH] Support importing chat completions settings --- crates/dump/src/reader/mod.rs | 9 +++++++++ crates/dump/src/reader/v6/mod.rs | 26 ++++++++++++++++++++++++++ crates/meilisearch/src/lib.rs | 28 +++++++++++++++++----------- 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/crates/dump/src/reader/mod.rs b/crates/dump/src/reader/mod.rs index 2b4440ab7..23e7eec9e 100644 --- a/crates/dump/src/reader/mod.rs +++ b/crates/dump/src/reader/mod.rs @@ -116,6 +116,15 @@ impl DumpReader { } } + pub fn chat_completions_settings( + &mut self, + ) -> Result> + '_>> { + match self { + DumpReader::Current(current) => current.chat_completions_settings(), + DumpReader::Compat(_compat) => Ok(Box::new(std::iter::empty())), + } + } + pub fn features(&self) -> Result> { match self { DumpReader::Current(current) => Ok(current.features()), diff --git a/crates/dump/src/reader/v6/mod.rs b/crates/dump/src/reader/v6/mod.rs index 0b4ba5bdd..0c920aadb 100644 --- a/crates/dump/src/reader/v6/mod.rs +++ b/crates/dump/src/reader/v6/mod.rs @@ -1,3 +1,4 @@ +use std::ffi::OsStr; use std::fs::{self, File}; use std::io::{BufRead, BufReader, ErrorKind}; use std::path::Path; @@ -21,6 +22,7 @@ pub type Unchecked = meilisearch_types::settings::Unchecked; pub type Task = crate::TaskDump; pub type Batch = meilisearch_types::batches::Batch; pub type Key = meilisearch_types::keys::Key; +pub type ChatCompletionSettings = meilisearch_types::features::ChatCompletionSettings; pub type RuntimeTogglableFeatures = meilisearch_types::features::RuntimeTogglableFeatures; pub type Network = meilisearch_types::features::Network; @@ -192,6 +194,30 @@ impl V6Reader { ) } + pub fn chat_completions_settings( + &mut self, + ) -> Result> + '_>> { + let entries = fs::read_dir(self.dump.path().join("chat-completions-settings"))?; + Ok(Box::new( + entries + .map(|entry| -> Result> { + let entry = entry?; + let file_name = entry.file_name(); + let path = Path::new(&file_name); + if entry.file_type()?.is_file() && path.extension() == Some(OsStr::new("json")) + { + let name = path.file_stem().unwrap().to_str().unwrap().to_string(); + let file = File::open(entry.path())?; + let settings = serde_json::from_reader(file)?; + Ok(Some((name, settings))) + } else { + Ok(None) + } + }) + .filter_map(|entry| entry.transpose()), + )) + } + pub fn features(&self) -> Option { self.features } diff --git a/crates/meilisearch/src/lib.rs b/crates/meilisearch/src/lib.rs index 871bd688e..b11a4a76d 100644 --- a/crates/meilisearch/src/lib.rs +++ b/crates/meilisearch/src/lib.rs @@ -498,14 +498,20 @@ fn import_dump( keys.push(key); } - // 3. Import the runtime features and network + // 3. Import the `ChatCompletionSettings`s. + for result in dump_reader.chat_completions_settings()? { + let (name, settings) = result?; + index_scheduler.put_chat_settings(&name, &settings)?; + } + + // 4. Import the runtime features and network let features = dump_reader.features()?.unwrap_or_default(); index_scheduler.put_runtime_features(features)?; let network = dump_reader.network()?.cloned().unwrap_or_default(); index_scheduler.put_network(network)?; - // 3.1 Use all cpus to process dump if `max_indexing_threads` not configured + // 4.1 Use all cpus to process dump if `max_indexing_threads` not configured let backup_config; let base_config = index_scheduler.indexer_config(); @@ -522,7 +528,7 @@ fn import_dump( // /!\ The tasks must be imported AFTER importing the indexes or else the scheduler might // try to process tasks while we're trying to import the indexes. - // 4. Import the indexes. + // 5. Import the indexes. for index_reader in dump_reader.indexes()? { let mut index_reader = index_reader?; let metadata = index_reader.metadata(); @@ -535,20 +541,20 @@ fn import_dump( let mut wtxn = index.write_txn()?; let mut builder = milli::update::Settings::new(&mut wtxn, &index, indexer_config); - // 4.1 Import the primary key if there is one. + // 5.1 Import the primary key if there is one. if let Some(ref primary_key) = metadata.primary_key { builder.set_primary_key(primary_key.to_string()); } - // 4.2 Import the settings. + // 5.2 Import the settings. tracing::info!("Importing the settings."); let settings = index_reader.settings()?; apply_settings_to_builder(&settings, &mut builder); let embedder_stats: Arc = Default::default(); builder.execute(&|| false, &progress, embedder_stats.clone())?; - // 4.3 Import the documents. - // 4.3.1 We need to recreate the grenad+obkv format accepted by the index. + // 5.3 Import the documents. + // 5.3.1 We need to recreate the grenad+obkv format accepted by the index. tracing::info!("Importing the documents."); let file = tempfile::tempfile()?; let mut builder = DocumentsBatchBuilder::new(BufWriter::new(file)); @@ -559,7 +565,7 @@ fn import_dump( // This flush the content of the batch builder. let file = builder.into_inner()?.into_inner()?; - // 4.3.2 We feed it to the milli index. + // 5.3.2 We feed it to the milli index. let reader = BufReader::new(file); let reader = DocumentsBatchReader::from_reader(reader)?; @@ -591,15 +597,15 @@ fn import_dump( index_scheduler.refresh_index_stats(&uid)?; } - // 5. Import the queue + // 6. Import the queue let mut index_scheduler_dump = index_scheduler.register_dumped_task()?; - // 5.1. Import the batches + // 6.1. Import the batches for ret in dump_reader.batches()? { let batch = ret?; index_scheduler_dump.register_dumped_batch(batch)?; } - // 5.2. Import the tasks + // 6.2. Import the tasks for ret in dump_reader.tasks()? { let (task, file) = ret?; index_scheduler_dump.register_dumped_task(task, file)?;