diff --git a/crates/cli/production/elasticsearch-cli/src/main.rs b/crates/cli/production/elasticsearch-cli/src/main.rs index 3a67c19b44..9c760bfbdb 100644 --- a/crates/cli/production/elasticsearch-cli/src/main.rs +++ b/crates/cli/production/elasticsearch-cli/src/main.rs @@ -6,21 +6,21 @@ //! The intent is to be able to quickly populate documents and indices. //! +use elasticsearch::http::transport::Transport; +use elasticsearch::Elasticsearch; +use log::info; +use sqlx::mysql::MySqlPoolOptions; +use sqlx::{MySql, Pool}; use std::collections::HashSet; use std::iter::FromIterator; -use elasticsearch::Elasticsearch; -use elasticsearch::http::transport::Transport; -use log::info; -use sqlx::{MySql, Pool}; -use sqlx::mysql::MySqlPoolOptions; use config::shared_constants::DEFAULT_RUST_LOG; -use elasticsearch_schema::searches::search_model_weights::{search_model_weights, SearchArgs}; +use elasticsearch_schema::searches::search_model_weights::{search_model_weights, ModelWeightsSortDirection, ModelWeightsSortField, SearchArgs}; use elasticsearch_schema::searches::search_tts_models::search_tts_models; use enums::by_table::model_weights::weights_types::WeightsType; use errors::AnyhowResult; -use crate::cli_args::{Action, Environment, parse_cli_args}; +use crate::cli_args::{parse_cli_args, Action, Environment}; use crate::plans::create_all_model_weight_documents::create_all_model_weight_documents; use crate::plans::create_all_tts_documents::create_all_tts_documents; use crate::plans::media_files::create_dimensional_media_file_documents::create_dimensional_media_file_documents; @@ -63,11 +63,15 @@ pub async fn main() -> AnyhowResult<()> { info!("Searching model weights..."); let results = search_model_weights(SearchArgs { - search_term: "zel", + //search_term: "zel", + search_term: "mariano", maybe_creator_user_token: None, maybe_ietf_primary_language_subtag: None, maybe_weights_categories: None, - maybe_weights_types: Some(HashSet::from_iter(vec![WeightsType::RvcV2])), + maybe_weights_types: Some(HashSet::from_iter(vec![WeightsType::Tacotron2])), + sort_field: Some(ModelWeightsSortField::UsageCount), + sort_direction: Some(ModelWeightsSortDirection::Ascending), + minimum_score: None, client: &elasticsearch, }).await?; diff --git a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs index b81a04c98a..9c56f75171 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -2,14 +2,17 @@ use std::collections::HashSet; use elasticsearch::{Elasticsearch, SearchParts}; use once_cell::sync::Lazy; -use serde_json::{json, Value}; +use serde_json::{json, Number, Value}; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; -use errors::{anyhow, AnyhowResult}; +use errors::AnyhowResult; use tokens::tokens::users::UserToken; -use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument}; +use crate::documents::model_weight_document::{ModelWeightDocument, MODEL_WEIGHT_INDEX}; + +/// Cut off scores below this threshold +const DEFAULT_MINIMUM_SCORE : u64 = 30; static JSON_QUERY : Lazy = Lazy::new(|| { const QUERY_TEMPLATE : &str = include_str!("../../../../../_elasticsearch/searches/model_weights/search.json"); @@ -20,99 +23,34 @@ static JSON_QUERY : Lazy = Lazy::new(|| { json }); -pub struct SearchModelWeightsQuery<'a> { - pub search_term: &'a str, - pub maybe_language_subtag: Option<&'a str>, - pub maybe_weights_type: Option, - pub maybe_weights_category: Option, +#[derive(Default, Copy, Clone, Eq, PartialEq)] +pub enum ModelWeightsSortField { + /// Sort based on the match score of the search term alone. + #[default] + MatchScore, + /// Sort based on the creation date + CreatedAt, + /// Sort based on the model usage count + UsageCount, + /// Sort based on the model bookmark count + BookmarkCount, + /// Sort based on the model positive ratings count + PositiveRatingCount, } -impl <'a>SearchModelWeightsQuery<'a> { +#[derive(Default, Copy, Clone, Eq, PartialEq)] +pub enum ModelWeightsSortDirection { + Ascending, + #[default] + Descending, +} - fn base_query(&self) -> Value { - json!({ - "query": { - "bool": { - "must": [ - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": self.search_term, - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": self.search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": self.search_term, - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - }) - } - - pub fn query(&self) -> AnyhowResult { - let mut search_json = self.base_query(); - - let must_clause = search_json.pointer_mut("/query/bool/must") - .map(|pointer| pointer.as_array_mut()) - .flatten() - .ok_or_else(|| anyhow!("could not get pointer to must clause"))?; - - if let Some(maybe_language_subtag) = &self.maybe_language_subtag { - must_clause.push(json!( - { - "match": { - "maybe_ietf_primary_language_subtag": maybe_language_subtag - } - } - )); +impl ModelWeightsSortDirection { + pub fn to_str(&self) -> &str { + match self { + Self::Ascending => "asc", + Self::Descending => "desc", } - - if let Some(weights_type) = &self.maybe_weights_type { - must_clause.push(json!( - { - "match": { - "weights_type": weights_type.to_string(), - } - } - )); - } - - if let Some(weights_category) = &self.maybe_weights_category { - must_clause.push(json!( - { - "match": { - "weights_category": weights_category.to_string(), - } - } - )); - } - - Ok(search_json) } } @@ -124,12 +62,21 @@ pub struct SearchArgs<'a> { pub maybe_weights_categories: Option>, pub maybe_weights_types: Option>, + pub sort_field: Option, + pub sort_direction: Option, + + pub minimum_score: Option, + pub client: &'a Elasticsearch, } pub async fn search_model_weights(args: SearchArgs<'_>) -> AnyhowResult> { let query = build_query(&args)?; + let json_query = serde_json::to_string(&query)?; + + println!("Query: {:#?}", json_query); + let search_response = args.client .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX])) .body(query) @@ -213,9 +160,71 @@ fn build_query(args: &SearchArgs) -> AnyhowResult { Some(json!(args.search_term)) })?; + let mut query = add_sort(query, args.sort_field, args.sort_direction); + + if let Some(mut object) = query.as_object_mut() { + object.insert( + "min_score".to_string(), + Value::Number(Number::from(args.minimum_score.unwrap_or(DEFAULT_MINIMUM_SCORE))) + ); + } + Ok(query) } +fn add_sort( + mut query: Value, + sort_field: Option, + sort_direction: Option, +) -> Value { + let sort_field = sort_field.unwrap_or_default(); + let sort_direction = sort_direction.unwrap_or_default(); + + if sort_field == ModelWeightsSortField::MatchScore + && sort_direction == ModelWeightsSortDirection::Descending { + return query; + } + + let sort = match sort_field { + ModelWeightsSortField::CreatedAt => json!([ + { + "created_at": { + "order": sort_direction.to_str(), + } + }, + "_score" + ]), + ModelWeightsSortField::UsageCount => json!([ + { + "cached_usage_count": { + "order": sort_direction.to_str(), + } + }, + "_score" + ]), + ModelWeightsSortField::BookmarkCount => json!([ + { + "bookmark_count": { + "order": sort_direction.to_str(), + } + }, + "_score" + ]), + ModelWeightsSortField::PositiveRatingCount => json!({ + "ratings_positive_count": { + "order": sort_direction.to_str(), + } + }), + _ => return query, + }; + + if let Some(mut object) = query.as_object_mut() { + object.insert("sort".to_string(), sort); + } + + query +} + fn must_be_not_deleted() -> Value { json!({ "term": { @@ -305,6 +314,9 @@ mod tests { maybe_ietf_primary_language_subtag: None, maybe_weights_categories: None, maybe_weights_types: None, + sort_field: Default::default(), + sort_direction: Default::default(), + minimum_score: None, client: &elasticsearch::Elasticsearch::default(), }).unwrap(); @@ -337,6 +349,9 @@ mod tests { maybe_ietf_primary_language_subtag: None, maybe_weights_categories: None, maybe_weights_types: None, + sort_field: Default::default(), + sort_direction: Default::default(), + minimum_score: None, client: &elasticsearch::Elasticsearch::default(), }).unwrap(); @@ -375,6 +390,9 @@ mod tests { maybe_weights_categories: None, maybe_weights_types: None, client: &elasticsearch::Elasticsearch::default(), + minimum_score: None, + sort_field: Default::default(), + sort_direction: Default::default(), }).unwrap(); let value = jsonpath_lib::select( @@ -414,6 +432,9 @@ mod tests { WeightsType::Tacotron2, WeightsType::GptSoVits, ])), + sort_field: Default::default(), + sort_direction: Default::default(), + minimum_score: None, client: &elasticsearch::Elasticsearch::default(), }).unwrap(); @@ -441,6 +462,9 @@ mod tests { WeightsCategory::VoiceConversion, ])), client: &elasticsearch::Elasticsearch::default(), + minimum_score: None, + sort_field: Default::default(), + sort_direction: Default::default(), }).unwrap(); let value = select(&search, "$.query.bool.must[0].bool.must[0].term.is_deleted"); @@ -469,165 +493,4 @@ mod tests { }) .collect() } - - mod old_tests { - use regex::Regex; - - use enums::by_table::model_weights::weights_category::WeightsCategory; - use enums::by_table::model_weights::weights_types::WeightsType; - use errors::AnyhowResult; - - use crate::searches::search_model_weights::SearchModelWeightsQuery; - - fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult { - let json = query.query()?; - let json = serde_json::to_string(&json)?; - Ok(json) - } - - fn compact_json(json: &str) -> String { - let regex = Regex::new("\\s+").expect("regex should parse"); - let json = regex.replace_all(&json, ""); - json.to_string() - } - - #[test] - fn default_search_only_keyword() { - let query = SearchModelWeightsQuery { - search_term: "FOO_BAR_BAZ", - maybe_language_subtag: None, - maybe_weights_type: None, - maybe_weights_category: None, - }; - - let json = query_to_json_string(query).unwrap(); - - // NB: Keys in emitted JSON are sorted. - let expected_json = r#" - { - "query": { - "bool": { - "must": [ - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": "FOO_BAR_BAZ", - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": "FOO_BAR_BAZ", - "boost": 1 - } - } - }, - { - "multi_match": { - "query": "FOO_BAR_BAZ", - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - } - "#; - - let expected_json = compact_json(&expected_json); - - assert_eq!(&json, &expected_json); - } - - #[test] - fn search_with_language_and_weights() { - let query = SearchModelWeightsQuery { - search_term: "FOO_BAR_BAZ", - maybe_language_subtag: Some("en"), - maybe_weights_type: Some(WeightsType::SoVitsSvc), - maybe_weights_category: Some(WeightsCategory::Vocoder), - }; - - let json = query_to_json_string(query).unwrap(); - - // NB: Keys in emitted JSON are sorted. - let expected_json = r#" - { - "query": { - "bool": { - "must": [ - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": "FOO_BAR_BAZ", - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": "FOO_BAR_BAZ", - "boost": 1 - } - } - }, - { - "multi_match": { - "query": "FOO_BAR_BAZ", - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - }, - { - "match": { - "maybe_ietf_primary_language_subtag": "en" - } - }, - { - "match": { - "weights_type": "so_vits_svc" - } - }, - { - "match": { - "weights_category": "vocoder" - } - } - ] - } - } - } - "#; - - let expected_json = compact_json(&expected_json); - - assert_eq!(&json, &expected_json); - } - } } diff --git a/crates/service/web/storyteller_web/src/docs/api_doc.rs b/crates/service/web/storyteller_web/src/docs/api_doc.rs index 8b6c99ee97..f531e7aca2 100644 --- a/crates/service/web/storyteller_web/src/docs/api_doc.rs +++ b/crates/service/web/storyteller_web/src/docs/api_doc.rs @@ -560,6 +560,8 @@ use crate::http_server::web_utils::response_success_helpers::*; SearchMediaFilesSuccessResponse, SearchModelWeightsError, SearchModelWeightsRequest, + SearchModelWeightsSortDirection, + SearchModelWeightsSortField, SearchModelWeightsSuccessResponse, SessionInfoError, SessionInfoSuccessResponse, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search/search_model_weights_impl.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search/search_model_weights_impl.rs index 6b6c10cc4d..761161d61f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search/search_model_weights_impl.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search/search_model_weights_impl.rs @@ -23,7 +23,7 @@ use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::state::server_state::ServerState; use crate::util::title_to_url_slug::title_to_url_slug; use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; -use elasticsearch_schema::searches::search_model_weights::{search_model_weights, SearchArgs}; +use elasticsearch_schema::searches::search_model_weights::{search_model_weights, ModelWeightsSortDirection, ModelWeightsSortField, SearchArgs}; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; @@ -37,6 +37,31 @@ pub struct SearchModelWeightsRequest { pub weight_type: Option, pub weight_category: Option, pub ietf_language_subtag: Option, + pub minimum_score: Option, + pub sort_field: Option, + pub sort_direction: Option, +} + +#[derive(Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum SearchModelWeightsSortField { + /// Sort based on the match score of the search term alone. + MatchScore, + /// Sort based on the creation date + CreatedAt, + /// Sort based on the model usage count + UsageCount, + /// Sort based on the model bookmark count + BookmarkCount, + /// Sort based on the model positive ratings count + PositiveRatingCount, +} + +#[derive(Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum SearchModelWeightsSortDirection { + Ascending, + Descending, } #[derive(Serialize, Clone, ToSchema)] @@ -138,12 +163,30 @@ pub async fn search_model_weights_impl( set }); + let sort_field = match request.sort_field { + Some(SearchModelWeightsSortField::MatchScore) => Some(ModelWeightsSortField::MatchScore), + Some(SearchModelWeightsSortField::CreatedAt) => Some(ModelWeightsSortField::CreatedAt), + Some(SearchModelWeightsSortField::UsageCount) => Some(ModelWeightsSortField::UsageCount), + Some(SearchModelWeightsSortField::BookmarkCount) => Some(ModelWeightsSortField::BookmarkCount), + Some(SearchModelWeightsSortField::PositiveRatingCount) => Some(ModelWeightsSortField::PositiveRatingCount), + None => None, + }; + + let sort_direction = match request.sort_direction { + Some(SearchModelWeightsSortDirection::Ascending) => Some(ModelWeightsSortDirection::Ascending), + Some(SearchModelWeightsSortDirection::Descending) => Some(ModelWeightsSortDirection::Descending), + None => None, + }; + let results = search_model_weights(SearchArgs { search_term: &request.search_term, maybe_creator_user_token: None, maybe_ietf_primary_language_subtag: request.ietf_language_subtag.as_deref(), maybe_weights_categories, maybe_weights_types, + sort_field, + sort_direction, + minimum_score: request.minimum_score, client: &server_state.elasticsearch, }) .await