From ac39a436d9cd5ab826145f37763b3e93786610a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Tue, 13 May 2025 16:33:58 +0200
Subject: [PATCH] Support overwriten prompts of the search query

---
 crates/index-scheduler/src/lib.rs     | 20 +++++++++++++++++---
 crates/meilisearch/src/routes/chat.rs | 20 +++++++++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/crates/index-scheduler/src/lib.rs b/crates/index-scheduler/src/lib.rs
index 4f1109348..75fab121c 100644
--- a/crates/index-scheduler/src/lib.rs
+++ b/crates/index-scheduler/src/lib.rs
@@ -53,8 +53,8 @@ use flate2::Compression;
 use meilisearch_types::batches::Batch;
 use meilisearch_types::features::{InstanceTogglableFeatures, Network, RuntimeTogglableFeatures};
 use meilisearch_types::heed::byteorder::BE;
-use meilisearch_types::heed::types::I128;
-use meilisearch_types::heed::{self, Env, RoTxn, WithoutTls};
+use meilisearch_types::heed::types::{Str, I128};
+use meilisearch_types::heed::{self, Database, Env, RoTxn, WithoutTls};
 use meilisearch_types::milli::index::IndexEmbeddingConfig;
 use meilisearch_types::milli::update::IndexerConfig;
 use meilisearch_types::milli::vector::{Embedder, EmbedderOptions, EmbeddingConfigs};
@@ -151,6 +151,9 @@ pub struct IndexScheduler {
     /// In charge of fetching and setting the status of experimental features.
     features: features::FeatureData,
 
+    /// Stores the custom prompts for the chat
+    chat_prompts: Database<Str, Str>,
+
     /// Everything related to the processing of the tasks
     pub scheduler: scheduler::Scheduler,
 
@@ -209,11 +212,16 @@ impl IndexScheduler {
             #[cfg(test)]
             run_loop_iteration: self.run_loop_iteration.clone(),
             features: self.features.clone(),
+            chat_prompts: self.chat_prompts.clone(),
         }
     }
 
     pub(crate) const fn nb_db() -> u32 {
-        Versioning::nb_db() + Queue::nb_db() + IndexMapper::nb_db() + features::FeatureData::nb_db()
+        Versioning::nb_db()
+            + Queue::nb_db()
+            + IndexMapper::nb_db()
+            + features::FeatureData::nb_db()
+            + 1 // chat-prompts
     }
 
     /// Create an index scheduler and start its run loop.
@@ -267,6 +275,7 @@ impl IndexScheduler {
         let features = features::FeatureData::new(&env, &mut wtxn, options.instance_features)?;
         let queue = Queue::new(&env, &mut wtxn, &options)?;
         let index_mapper = IndexMapper::new(&env, &mut wtxn, &options, budget)?;
+        let chat_prompts = env.create_database(&mut wtxn, Some("chat-prompts"))?;
         wtxn.commit()?;
 
         // allow unreachable_code to get rids of the warning in the case of a test build.
@@ -290,6 +299,7 @@ impl IndexScheduler {
             #[cfg(test)]
             run_loop_iteration: Arc::new(RwLock::new(0)),
             features,
+            chat_prompts,
         };
 
         this.run();
@@ -857,6 +867,10 @@ impl IndexScheduler {
             .collect();
         res.map(EmbeddingConfigs::new)
     }
+
+    pub fn chat_prompts<'t>(&self, rtxn: &'t RoTxn, name: &str) -> heed::Result<Option<&'t str>> {
+        self.chat_prompts.get(rtxn, name)
+    }
 }
 
 /// The outcome of calling the [`IndexScheduler::tick`] function.
diff --git a/crates/meilisearch/src/routes/chat.rs b/crates/meilisearch/src/routes/chat.rs
index 335d5bf43..3e5bf6957 100644
--- a/crates/meilisearch/src/routes/chat.rs
+++ b/crates/meilisearch/src/routes/chat.rs
@@ -17,7 +17,7 @@ use meilisearch_types::milli::index::IndexEmbeddingConfig;
 use meilisearch_types::milli::prompt::PromptData;
 use meilisearch_types::milli::vector::EmbeddingConfig;
 use meilisearch_types::{Document, Index};
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
 
 use crate::extractors::authentication::policies::ActionPolicy;
@@ -68,6 +68,24 @@ async fn chat(
         "Meilisearch /chat only support one completion at a time (n = 1, n = null)"
     );
 
+    let rtxn = index_scheduler.read_txn().unwrap();
+    let search_in_index_description = index_scheduler
+        .chat_prompts(&rtxn, "searchInIndex-description")
+        .unwrap()
+        .unwrap_or(DEFAULT_SEARCH_IN_INDEX_TOOL_DESCRIPTION)
+        .to_string();
+    let search_in_index_q_param_description = index_scheduler
+        .chat_prompts(&rtxn, "searchInIndex-q-param-description")
+        .unwrap()
+        .unwrap_or(DEFAULT_SEARCH_IN_INDEX_Q_PARAMETER_TOOL_DESCRIPTION)
+        .to_string();
+    let search_in_index_index_description = index_scheduler
+        .chat_prompts(&rtxn, "searchInIndex-index-param-description")
+        .unwrap()
+        .unwrap_or(DEFAULT_SEARCH_IN_INDEX_INDEX_PARAMETER_TOOL_DESCRIPTION)
+        .to_string();
+    drop(rtxn);
+
     let mut response;
     loop {
         let mut tools = chat_completion.tools.get_or_insert_default();