diff --git a/Cargo.lock b/Cargo.lock index 7ef5187f6..e6f4d9a6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3592,6 +3592,7 @@ dependencies = [ "brotli 6.0.0", "bstr", "build-info", + "bumpalo", "byte-unit", "bytes", "cargo_toml", diff --git a/crates/meilisearch/Cargo.toml b/crates/meilisearch/Cargo.toml index 398b62dad..62f7cfa0a 100644 --- a/crates/meilisearch/Cargo.toml +++ b/crates/meilisearch/Cargo.toml @@ -32,6 +32,7 @@ async-trait = "0.1.85" bstr = "1.11.3" byte-unit = { version = "5.1.6", features = ["serde"] } bytes = "1.9.0" +bumpalo = "3.16.0" clap = { version = "4.5.24", features = ["derive", "env"] } crossbeam-channel = "0.5.15" deserr = { version = "0.6.3", features = ["actix-web"] } diff --git a/crates/meilisearch/src/routes/chat.rs b/crates/meilisearch/src/routes/chat.rs index d85c14c36..c34bbcaea 100644 --- a/crates/meilisearch/src/routes/chat.rs +++ b/crates/meilisearch/src/routes/chat.rs @@ -1,5 +1,7 @@ +use std::cell::RefCell; use std::collections::HashMap; use std::mem; +use std::sync::RwLock; use std::time::Duration; use actix_web::web::{self, Data}; @@ -16,27 +18,33 @@ use async_openai::types::{ FunctionObjectArgs, }; use async_openai::Client; +use bumpalo::Bump; use futures::StreamExt; use index_scheduler::IndexScheduler; use meilisearch_auth::AuthController; use meilisearch_types::error::ResponseError; +use meilisearch_types::heed::RoTxn; use meilisearch_types::keys::actions; -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, Serialize}; +use meilisearch_types::milli::index::ChatConfig; +use meilisearch_types::milli::prompt::{Prompt, PromptData}; +use meilisearch_types::milli::update::new::document::DocumentFromDb; +use meilisearch_types::milli::{ + DocumentId, FieldIdMapWithMetadata, GlobalFieldsIdsMap, MetadataBuilder, TimeBudget, +}; +use meilisearch_types::Index; +use serde::Deserialize; use serde_json::json; use tokio::runtime::Handle; use tokio::sync::mpsc::error::SendError; use super::settings::chat::{ChatPrompts, GlobalChatSettings}; +use crate::error::MeilisearchHttpError; use crate::extractors::authentication::policies::ActionPolicy; use crate::extractors::authentication::{extract_token_from_request, GuardedData, Policy as _}; use crate::metrics::MEILISEARCH_DEGRADED_SEARCH_REQUESTS; use crate::routes::indexes::search::search_kind; use crate::search::{ - add_search_rules, perform_search, HybridQuery, RetrieveVectors, SearchQuery, SemanticRatio, + add_search_rules, prepare_search, search_from_kind, HybridQuery, SearchQuery, SemanticRatio, }; use crate::search_queue::SearchQueue; @@ -175,15 +183,22 @@ async fn process_search_request( let permit = search_queue.try_get_search_permit().await?; let features = index_scheduler.features(); let index_cloned = index.clone(); - let search_result = tokio::task::spawn_blocking(move || { - perform_search( - index_uid.to_string(), - &index_cloned, - query, - search_kind, - RetrieveVectors::new(false), - features, - ) + let search_result = tokio::task::spawn_blocking(move || -> Result<_, ResponseError> { + let rtxn = index_cloned.read_txn()?; + let time_budget = match index_cloned + .search_cutoff(&rtxn) + .map_err(|e| MeilisearchHttpError::from_milli(e, Some(index_uid.clone())))? + { + Some(cutoff) => TimeBudget::new(Duration::from_millis(cutoff)), + None => TimeBudget::default(), + }; + + 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, _)| search_results) + .map_err(ResponseError::from) }) .await; permit.drop().await; @@ -198,9 +213,11 @@ async fn process_search_request( // analytics.publish(aggregate, &req); let search_result = search_result?; - let formatted = - format_documents(&index, search_result.hits.into_iter().map(|doc| doc.document)); + let rtxn = index.read_txn()?; + let render_alloc = Bump::new(); + let formatted = format_documents(&rtxn, &index, &render_alloc, search_result.documents_ids)?; let text = formatted.join("\n"); + drop(rtxn); Ok((index, text)) } @@ -506,31 +523,36 @@ struct SearchInIndexParameters { q: Option, } -fn format_documents(index: &Index, documents: impl Iterator) -> Vec { - let rtxn = index.read_txn().unwrap(); - let IndexEmbeddingConfig { name: _, config, user_provided: _ } = index - .embedding_configs(&rtxn) - .unwrap() +fn format_documents<'t, 'doc>( + rtxn: &RoTxn<'t>, + index: &Index, + doc_alloc: &'doc Bump, + internal_docids: Vec, +) -> Result, ResponseError> { + let ChatConfig { prompt: PromptData { template, max_bytes }, .. } = index.chat_config(rtxn)?; + + let prompt = Prompt::new(template, max_bytes).unwrap(); + let fid_map = index.fields_ids_map(rtxn)?; + let metadata_builder = MetadataBuilder::from_index(index, rtxn)?; + let fid_map_with_meta = FieldIdMapWithMetadata::new(fid_map.clone(), metadata_builder); + let global = RwLock::new(fid_map_with_meta); + let gfid_map = RefCell::new(GlobalFieldsIdsMap::new(&global)); + + let external_ids: Vec = index + .external_id_of(rtxn, internal_docids.iter().copied())? .into_iter() - .find(|conf| conf.name == EMBEDDER_NAME) - .unwrap(); + .collect::>()?; - let EmbeddingConfig { - embedder_options: _, - prompt: PromptData { template, max_bytes: _ }, - quantized: _, - } = config; + let mut renders = Vec::new(); + for (docid, external_docid) in internal_docids.into_iter().zip(external_ids) { + let document = match DocumentFromDb::new(docid, rtxn, index, &fid_map)? { + Some(doc) => doc, + None => continue, + }; - #[derive(Serialize)] - struct Doc { - doc: T, + let text = prompt.render_document(&external_docid, document, &gfid_map, doc_alloc).unwrap(); + renders.push(text); } - let template = liquid::ParserBuilder::with_stdlib().build().unwrap().parse(&template).unwrap(); - documents - .map(|doc| { - let object = liquid::to_object(&Doc { doc }).unwrap(); - template.render(&object).unwrap() - }) - .collect() + Ok(renders) } diff --git a/crates/meilisearch/src/routes/settings/chat.rs b/crates/meilisearch/src/routes/settings/chat.rs index a971ad102..586fa041e 100644 --- a/crates/meilisearch/src/routes/settings/chat.rs +++ b/crates/meilisearch/src/routes/settings/chat.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use actix_web::web::{self, Data}; use actix_web::HttpResponse; use index_scheduler::IndexScheduler; @@ -51,7 +49,6 @@ pub struct GlobalChatSettings { pub base_api: Option, pub api_key: Option, pub prompts: ChatPrompts, - pub indexes: BTreeMap, } #[derive(Debug, Serialize, Deserialize)] @@ -105,7 +102,6 @@ impl Default for GlobalChatSettings { .to_string(), pre_query: "".to_string(), }, - indexes: BTreeMap::new(), } } } diff --git a/crates/meilisearch/src/search/mod.rs b/crates/meilisearch/src/search/mod.rs index 1dd16c474..16d04cd58 100644 --- a/crates/meilisearch/src/search/mod.rs +++ b/crates/meilisearch/src/search/mod.rs @@ -882,7 +882,7 @@ pub fn add_search_rules(filter: &mut Option, rules: IndexSearchRules) { } } -fn prepare_search<'t>( +pub fn prepare_search<'t>( index: &'t Index, rtxn: &'t RoTxn, query: &'t SearchQuery, diff --git a/crates/milli/src/external_documents_ids.rs b/crates/milli/src/external_documents_ids.rs index 755b801ec..598465e5f 100644 --- a/crates/milli/src/external_documents_ids.rs +++ b/crates/milli/src/external_documents_ids.rs @@ -32,13 +32,13 @@ impl ExternalDocumentsIds { &self, rtxn: &RoTxn<'_>, external_id: A, - ) -> heed::Result> { + ) -> heed::Result> { self.0.get(rtxn, external_id.as_ref()) } /// An helper function to debug this type, returns an `HashMap` of both, /// soft and hard fst maps, combined. - pub fn to_hash_map(&self, rtxn: &RoTxn<'_>) -> heed::Result> { + pub fn to_hash_map(&self, rtxn: &RoTxn<'_>) -> heed::Result> { let mut map = HashMap::default(); for result in self.0.iter(rtxn)? { let (external, internal) = result?; diff --git a/crates/milli/src/fields_ids_map.rs b/crates/milli/src/fields_ids_map.rs index 9a016e7bd..d2abd840c 100644 --- a/crates/milli/src/fields_ids_map.rs +++ b/crates/milli/src/fields_ids_map.rs @@ -7,6 +7,7 @@ use crate::FieldId; mod global; pub mod metadata; pub use global::GlobalFieldsIdsMap; +pub use metadata::{FieldIdMapWithMetadata, MetadataBuilder}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FieldsIdsMap { diff --git a/crates/milli/src/lib.rs b/crates/milli/src/lib.rs index 47d3dc75c..504b4c68d 100644 --- a/crates/milli/src/lib.rs +++ b/crates/milli/src/lib.rs @@ -52,18 +52,19 @@ pub use search::new::{ }; use serde_json::Value; pub use thread_pool_no_abort::{PanicCatched, ThreadPoolNoAbort, ThreadPoolNoAbortBuilder}; -pub use {charabia as tokenizer, heed, rhai}; +pub use {arroy, charabia as tokenizer, heed, rhai}; pub use self::asc_desc::{AscDesc, AscDescError, Member, SortError}; -pub use self::attribute_patterns::AttributePatterns; -pub use self::attribute_patterns::PatternMatch; +pub use self::attribute_patterns::{AttributePatterns, PatternMatch}; pub use self::criterion::{default_criteria, Criterion, CriterionError}; pub use self::error::{ Error, FieldIdMapMissingEntry, InternalError, SerializationError, UserError, }; pub use self::external_documents_ids::ExternalDocumentsIds; pub use self::fieldids_weights_map::FieldidsWeightsMap; -pub use self::fields_ids_map::{FieldsIdsMap, GlobalFieldsIdsMap}; +pub use self::fields_ids_map::{ + FieldIdMapWithMetadata, FieldsIdsMap, GlobalFieldsIdsMap, MetadataBuilder, +}; pub use self::filterable_attributes_rules::{ FilterFeatures, FilterableAttributesFeatures, FilterableAttributesPatterns, FilterableAttributesRule, @@ -84,8 +85,6 @@ pub use self::search::{ }; pub use self::update::ChannelCongestion; -pub use arroy; - pub type Result = std::result::Result; pub type Attribute = u32;