Add the index descriptions to the function description

This commit is contained in:
Clément Renault 2025-05-22 10:40:43 +02:00 committed by Kerollmops
parent 7bfb7b7d58
commit 1e67a47a8f
No known key found for this signature in database
GPG Key ID: F250A4C4E3AE5F5F

View File

@ -1,5 +1,6 @@
use std::cell::RefCell; use std::cell::RefCell;
use std::collections::HashMap; use std::collections::HashMap;
use std::fmt::Write as _;
use std::mem; use std::mem;
use std::sync::RwLock; use std::sync::RwLock;
use std::time::Duration; use std::time::Duration;
@ -97,18 +98,29 @@ fn setup_search_tool(
panic!("{SEARCH_IN_INDEX_FUNCTION_NAME} function already set"); panic!("{SEARCH_IN_INDEX_FUNCTION_NAME} function already set");
} }
let index_uids: Vec<_> = index_scheduler let mut index_uids = Vec::new();
.index_names()? let mut function_description = prompts.search_description.clone().unwrap();
.into_iter() index_scheduler.try_for_each_index::<_, ()>(|name, index| {
.filter(|index_uid| filters.is_index_authorized(&index_uid)) // Make sure to skip unauthorized indexes
.collect(); if !filters.is_index_authorized(&name) {
return Ok(());
}
let rtxn = index.read_txn()?;
let chat_config = index.chat_config(&rtxn)?;
let index_description = chat_config.description;
let _ = writeln!(&mut function_description, "\n\n - {name}: {index_description}\n");
index_uids.push(name.to_string());
Ok(())
})?;
let tool = ChatCompletionToolArgs::default() let tool = ChatCompletionToolArgs::default()
.r#type(ChatCompletionToolType::Function) .r#type(ChatCompletionToolType::Function)
.function( .function(
FunctionObjectArgs::default() FunctionObjectArgs::default()
.name(SEARCH_IN_INDEX_FUNCTION_NAME) .name(SEARCH_IN_INDEX_FUNCTION_NAME)
.description(&prompts.search_description.clone().unwrap()) .description(&function_description)
.parameters(json!({ .parameters(json!({
"type": "object", "type": "object",
"properties": { "properties": {