Merge 08e096c985 into fef089c7b6

crates/meilisearch/src/routes/chats/chat_completions.rs

@@ -27,9 +27,10 @@ use meilisearch_types::features::{
     ChatCompletionPrompts as DbChatCompletionPrompts,
     ChatCompletionSource as DbChatCompletionSource, SystemRole,
 };
+use meilisearch_types::heed::RoTxn;
 use meilisearch_types::keys::actions;
 use meilisearch_types::milli::index::ChatConfig;
-use meilisearch_types::milli::{all_obkv_to_json, obkv_to_json, TimeBudget};
+use meilisearch_types::milli::{all_obkv_to_json, obkv_to_json, OrderBy, TimeBudget};
 use meilisearch_types::{Document, Index};
 use serde::Deserialize;
 use serde_json::json;
@@ -165,6 +166,7 @@ fn setup_search_tool(
 
     let mut index_uids = Vec::new();
     let mut function_description = prompts.search_description.clone();
+    let mut filter_description = prompts.search_filter_param.clone();
     index_scheduler.try_for_each_index::<_, ()>(|name, index| {
         // Make sure to skip unauthorized indexes
         if !filters.is_index_authorized(name) {
@@ -176,16 +178,22 @@ fn setup_search_tool(
         let index_description = chat_config.description;
         let _ = writeln!(&mut function_description, "\n\n - {name}: {index_description}\n");
         index_uids.push(name.to_string());
+        let facet_distributions = format_facet_distributions(&index, &rtxn, 10).unwrap(); // TODO do not unwrap
+        let _ = writeln!(&mut filter_description, "\n## Facet distributions of the {name} index");
+        let _ = writeln!(&mut filter_description, "{facet_distributions}");
 
         Ok(())
     })?;
 
+    tracing::debug!("LLM function description: {function_description}");
+    tracing::debug!("LLM filter description: {filter_description}");
+
     let tool = ChatCompletionToolArgs::default()
         .r#type(ChatCompletionToolType::Function)
         .function(
             FunctionObjectArgs::default()
                 .name(MEILI_SEARCH_IN_INDEX_FUNCTION_NAME)
-                .description(&function_description)
+                .description(function_description)
                 .parameters(json!({
                     "type": "object",
                     "properties": {
@@ -199,9 +207,13 @@ fn setup_search_tool(
                             // "type": ["string", "null"],
                             "type": "string",
                             "description": prompts.search_q_param,
+                        },
+                        "filter": {
+                            "type": "string",
+                            "description": filter_description,
                         }
                     },
-                    "required": ["index_uid", "q"],
+                    "required": ["index_uid", "q", "filter"],
                     "additionalProperties": false,
                 }))
                 .strict(true)
@@ -243,11 +255,19 @@ async fn process_search_request(
     auth_token: &str,
     index_uid: String,
     q: Option<String>,
+    filter: Option<String>,
 ) -> Result<(Index, Vec<Document>, String), ResponseError> {
     let index = index_scheduler.index(&index_uid)?;
     let rtxn = index.static_read_txn()?;
     let ChatConfig { description: _, prompt: _, search_parameters } = index.chat_config(&rtxn)?;
-    let mut query = SearchQuery { q, ..SearchQuery::from(search_parameters) };
+    let mut query = SearchQuery {
+        q,
+        filter: filter.map(serde_json::Value::from),
+        ..SearchQuery::from(search_parameters)
+    };
+
+    tracing::debug!("LLM query: {:?}", query);
+
     let auth_filter = ActionPolicy::<{ actions::SEARCH }>::authenticate(
         auth_ctrl,
         auth_token,
@@ -276,14 +296,23 @@ async fn process_search_request(
         let (search, _is_finite_pagination, _max_total_hits, _offset) =
             prepare_search(&index_cloned, &rtxn, &query, &search_kind, time_budget, features)?;
 
-        search_from_kind(index_uid, search_kind, search)
-            .map(|(search_results, _)| (rtxn, search_results))
-            .map_err(ResponseError::from)
+        match search_from_kind(index_uid, search_kind, search) {
+            Ok((search_results, _)) => Ok((rtxn, Ok(search_results))),
+            Err(MeilisearchHttpError::Milli {
+                error: meilisearch_types::milli::Error::UserError(user_error),
+                index_name: _,
+            }) => Ok((rtxn, Err(user_error))),
+            Err(err) => Err(ResponseError::from(err)),
+        }
     })
     .await;
     permit.drop().await;
 
-    let output = output?;
+    let output = match output? {
+        Ok((rtxn, Ok(search_results))) => Ok((rtxn, search_results)),
+        Ok((_rtxn, Err(error))) => return Ok((index, Vec::new(), error.to_string())),
+        Err(err) => Err(ResponseError::from(err)),
+    };
     let mut documents = Vec::new();
     if let Ok((ref rtxn, ref search_result)) = output {
         // aggregate.succeed(search_result);
@@ -391,16 +420,19 @@ async fn non_streamed_chat(
 
                 for call in meili_calls {
                     let result = match serde_json::from_str(&call.function.arguments) {
-                        Ok(SearchInIndexParameters { index_uid, q }) => process_search_request(
-                            &index_scheduler,
-                            auth_ctrl.clone(),
-                            &search_queue,
-                            auth_token,
-                            index_uid,
-                            q,
-                        )
-                        .await
-                        .map_err(|e| e.to_string()),
+                        Ok(SearchInIndexParameters { index_uid, q, filter }) => {
+                            process_search_request(
+                                &index_scheduler,
+                                auth_ctrl.clone(),
+                                &search_queue,
+                                auth_token,
+                                index_uid,
+                                q,
+                                filter,
+                            )
+                            .await
+                            .map_err(|e| e.to_string())
+                        }
                         Err(err) => Err(err.to_string()),
                     };
 
@@ -691,13 +723,14 @@ async fn handle_meili_tools(
         let mut error = None;
 
         let result = match serde_json::from_str(&call.function.arguments) {
-            Ok(SearchInIndexParameters { index_uid, q }) => match process_search_request(
+            Ok(SearchInIndexParameters { index_uid, q, filter }) => match process_search_request(
                 index_scheduler,
                 auth_ctrl.clone(),
                 search_queue,
                 auth_token,
                 index_uid,
                 q,
+                filter,
             )
             .await
             {
@@ -773,4 +806,42 @@ struct SearchInIndexParameters {
     index_uid: String,
     /// The query parameter to use.
     q: Option<String>,
+    /// The filter parameter to use.
+    filter: Option<String>,
+}
+
+fn format_facet_distributions(
+    index: &Index,
+    rtxn: &RoTxn,
+    max_values_per_facet: usize,
+) -> meilisearch_types::milli::Result<String> {
+    let universe = index.documents_ids(&rtxn)?;
+    let rules = index.filterable_attributes_rules(&rtxn)?;
+    let fields_ids_map = index.fields_ids_map(&rtxn)?;
+    let filterable_attributes = fields_ids_map
+        .names()
+        .filter(|name| rules.iter().any(|rule| rule.match_str(name).matches()))
+        .map(|name| (name, OrderBy::Count));
+    let facets_distribution = index
+        .facets_distribution(&rtxn)
+        .max_values_per_facet(max_values_per_facet)
+        .candidates(universe)
+        .facets(filterable_attributes)
+        .execute()?;
+
+    let mut output = String::new();
+    for (facet_name, entries) in facets_distribution {
+        let _ = write!(&mut output, "{}: ", facet_name);
+        let total_entries = entries.len();
+        for (i, (value, count)) in entries.into_iter().enumerate() {
+            let _ = if total_entries.saturating_sub(1) == i {
+                write!(&mut output, "{} ({}).", value, count)
+            } else {
+                write!(&mut output, "{} ({}), ", value, count)
+            };
+        }
+        let _ = writeln!(&mut output);
+    }
+
+    Ok(output)
 }