From 35231e5f82892268d34458729382f2a3129ffd1c Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 17 Jan 2024 23:13:47 -0500 Subject: [PATCH 1/7] fix last test failures --- .../payloads/generic_inference_args/generic_inference_args.rs | 2 +- crates/lib/mysql_queries/src/queries/model_weights/tests.rs | 2 -- .../src/job/job_types/image_generation/sd/process_job.rs | 2 ++ 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index eae5b1de77..656df4fb13 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -237,7 +237,7 @@ mod tests { }; let json = serde_json::ser::to_string(&args).unwrap(); - assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"vs":"video_media_token","is":"image_media_token","sd":"sd_model_token","lm":"lora_model_token","p":"prompt","ap":"a_prompt","np":"n_prompt","se":1,"mu":"upload_path","lu":"lora_upload_path"}}}"#.to_string()); + assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"sd":"sd_model_token","lm":"lora_model_token","w":512,"h":512,"s":"sampler","p":"prompt","np":"n_prompt","se":1,"mu":"upload_path","cf":7,"lu":"lora_upload_path","sa":25,"bc":1,"t":"type_of_inference","de":"Option","na":"Model Name"}}}"#.to_string()); } #[test] diff --git a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs index 13a70c3072..7f72eb2fd5 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs @@ -293,8 +293,6 @@ mod tests { maybe_scoped_weight_type: None, maybe_scoped_weight_category: None, mysql_pool: &pool, - maybe_scoped_weight_type: None, - maybe_scoped_weight_category: None, }).await?; for weight in weights_by_username.records.iter() { // print weight diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs index b7ea33cf16..d020ea1faf 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs @@ -427,6 +427,7 @@ pub async fn process_job_inference( } +#[ignore] #[cfg(test)] mod tests { use std::path::PathBuf; @@ -439,6 +440,7 @@ mod tests { }; use errors::AnyhowResult; + #[ignore] #[tokio::test] async fn test_seed_weights_files() -> AnyhowResult<()> { let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data"); From c91f5488786d0f78db569f7be13701a24a02af8f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 18 Jan 2024 01:58:36 -0500 Subject: [PATCH 2/7] hack elasticsearch options --- crates/cli/elasticsearch-cli/src/main.rs | 6 +- .../src/searches/search_model_weights.rs | 65 ++++++++++++++++--- .../weights/search_model_weights_handler.rs | 6 +- 3 files changed, 66 insertions(+), 11 deletions(-) diff --git a/crates/cli/elasticsearch-cli/src/main.rs b/crates/cli/elasticsearch-cli/src/main.rs index 06afd4f47e..dbe8a179c2 100644 --- a/crates/cli/elasticsearch-cli/src/main.rs +++ b/crates/cli/elasticsearch-cli/src/main.rs @@ -14,6 +14,7 @@ use sqlx::mysql::MySqlPoolOptions; use config::shared_constants::DEFAULT_RUST_LOG; use elasticsearch_schema::searches::search_model_weights::search_model_weights; 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}; @@ -55,9 +56,10 @@ pub async fn main() -> AnyhowResult<()> { } Action::SearchModelWeights => { info!("Searching model weights..."); - let results = search_model_weights(&elasticsearch, "zel", Some("en")).await?; + //let results = search_model_weights(&elasticsearch, "zel", Some("en"), Some(WeightsType::SoVitsSvc)).await?; + let results = search_model_weights(&elasticsearch, "zel", None, Some(WeightsType::RvcV2)).await?; for result in results { - println!("Result: {:?}", result); + println!("Result: {:#?}", result); } } } 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 d6caad142f..85a5143744 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -1,6 +1,7 @@ use elasticsearch::{Elasticsearch, SearchParts}; use log::error; use serde_json::{json, Value}; +use enums::by_table::model_weights::weights_types::WeightsType; use errors::AnyhowResult; @@ -10,11 +11,13 @@ pub async fn search_model_weights( client: &Elasticsearch, search_term: &str, maybe_language_subtag: Option<&str>, + maybe_weights_type: Option, ) -> AnyhowResult> { - let search_json = match maybe_language_subtag { - None => query_model_weights(search_term), - Some(language) => query_model_weights_with_required_language(search_term, language) + let search_json = match (maybe_language_subtag, maybe_weights_type) { + (Some(language), None) => query_model_weights_with_required_language(search_term, language), + (None, Some(weights_type)) => query_model_weights_with_model_weights_type(search_term, weights_type), + _ => query_model_weights(search_term), }; let search_response = client @@ -44,11 +47,7 @@ pub async fn search_model_weights( let maybe_object = hit.get_mut("_source") .map(|source| source.take()); if let Some(value) = maybe_object { - let document = serde_json::from_value::(value) - .map_err(|err| { - error!("Error: {:?}", &err); - err - })?; + let document = serde_json::from_value::(value)?; documents.push(document); } } @@ -153,3 +152,53 @@ fn query_model_weights_with_required_language(search_term: &str, language_tag: & } }) } + +fn query_model_weights_with_model_weights_type(search_term: &str, weight_type: WeightsType) -> Value { + json!({ + "query": { + "bool": { + "must": [ + { + "match": { + "weights_type": weight_type.to_string(), + } + }, + { + "bool": { + "should": [ + { + "fuzzy": { + "title": { + "value": search_term, + "fuzziness": 2 + } + } + }, + { + "match": { + "title": { + "query": search_term, + "boost": 1 + } + } + }, + { + "multi_match": { + "query": search_term, + "type": "bool_prefix", + "fields": [ + "title", + "title._2gram", + "title._3gram" + ], + "boost": 50 + } + } + ] + } + } + ] + } + } + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index e27ece3566..aab82b79bd 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -30,6 +30,9 @@ use crate::server_state::ServerState; #[derive(Deserialize, ToSchema)] pub struct SearchModelWeightsRequest { pub search_term: String, + pub weight_type: Option, + pub weight_category: Option, + pub ietf_language_subtag: Option, } #[derive(Serialize, Clone, ToSchema)] @@ -113,7 +116,8 @@ pub async fn search_model_weights_handler( let results = search_model_weights( &server_state.elasticsearch, &request.search_term, - None) + request.ietf_language_subtag.as_deref(), + request.weight_type) .await .map_err(|err| { error!("Searching error: {:?}", err); From c15360ddb3ab6c837fe72107463a516658ec2fed Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 18 Jan 2024 18:17:21 -0500 Subject: [PATCH 3/7] ElasticSearch test --- Cargo.lock | 1 + crates/lib/elasticsearch_schema/Cargo.toml | 3 +- .../src/searches/search_model_weights.rs | 67 ++++++++++++++++++- 3 files changed, 68 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fcc378de3e..a2ac2e31bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2149,6 +2149,7 @@ dependencies = [ "errors", "log", "my-workspace-hack", + "regex", "serde", "serde_json", "tokens", diff --git a/crates/lib/elasticsearch_schema/Cargo.toml b/crates/lib/elasticsearch_schema/Cargo.toml index 37dac8d081..2f974bac32 100644 --- a/crates/lib/elasticsearch_schema/Cargo.toml +++ b/crates/lib/elasticsearch_schema/Cargo.toml @@ -27,6 +27,7 @@ chrono = { version = "0.4.22", features = ["serde"] } serde = { version = "1.0.152", features = ["derive"] } serde_json = "1.0.108" my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } +regex = "1.10.2" [dev-dependencies] -# None yet +# None yet \ No newline at end of file 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 85a5143744..1a2553e593 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -1,8 +1,7 @@ use elasticsearch::{Elasticsearch, SearchParts}; -use log::error; use serde_json::{json, Value}; -use enums::by_table::model_weights::weights_types::WeightsType; +use enums::by_table::model_weights::weights_types::WeightsType; use errors::AnyhowResult; use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument}; @@ -202,3 +201,67 @@ fn query_model_weights_with_model_weights_type(search_term: &str, weight_type: W } }) } + +#[cfg(test)] +mod tests { + use regex::Regex; + use crate::searches::search_model_weights::query_model_weights; + + #[test] + fn default_search_only_keyword() { + let json = query_model_weights("FOO_BAR_BAZ"); + let json = serde_json::to_string(&json).unwrap(); + + // NB: Keys in emitted JSON are sorted. + let expected_json = r#" + { + "query": { + "bool": { + "must": [ + { + "bool": { + "should": [ + { + "fuzzy": { + "title": { + "fuzziness": 2, + "value": "FOO_BAR_BAZ" + } + } + }, + { + "match": { + "title": { + "boost": 1, + "query": "FOO_BAR_BAZ" + } + } + }, + { + "multi_match": { + "boost": 50, + "fields": [ + "title", + "title._2gram", + "title._3gram" + ], + "query": "FOO_BAR_BAZ", + "type": "bool_prefix" + } + } + ] + } + } + ] + } + } + } + "#; + + let regex = Regex::new("\\s+").expect("regex should parse"); + + let expected_json = regex.replace_all(&expected_json, ""); + + assert_eq!(&json, &expected_json); + } +} From 3318067cb58c6c5e5e8f97fdc0c3b7ea7b3b09c3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 18 Jan 2024 19:24:39 -0500 Subject: [PATCH 4/7] hide deleted files from mods/admins - it's just clutter --- .../voice_dataset_samples/list_samples_by_dataset.rs | 8 ++++++-- .../voice_datasets/list_datasets_by_session.rs | 8 ++++++-- .../voice_datasets/list_datasets_by_user.rs | 8 ++++++-- .../voice_designer/voices/list_available_voices.rs | 6 ++++-- .../voice_designer/voices/list_voices_by_session.rs | 8 ++++++-- .../voice_designer/voices/list_voices_by_user.rs | 8 ++++++-- 6 files changed, 34 insertions(+), 12 deletions(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_dataset_samples/list_samples_by_dataset.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_dataset_samples/list_samples_by_dataset.rs index fed08dceec..da5d4500eb 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_dataset_samples/list_samples_by_dataset.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_dataset_samples/list_samples_by_dataset.rs @@ -86,11 +86,15 @@ pub async fn list_samples_by_dataset_handler( }; let dataset_token = ZsVoiceDatasetToken(path.dataset_token.clone()); - let is_mod = user_session.can_ban_users; + let _is_mod = user_session.can_ban_users; + + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + // We should focus on visibility-related controls instead. + const CAN_SEE_DELETED : bool = false; let query_results = list_dataset_samples_for_dataset_token( &dataset_token, - is_mod, + CAN_SEE_DELETED, &server_state.mysql_pool ).await.map_err(|e| { warn!("list_samples error: {:?}", e); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_session.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_session.rs index 8d7fa4e38f..f7e781cad1 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_session.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_session.rs @@ -77,12 +77,16 @@ pub async fn list_datasets_by_session_handler( })?; - let is_mod = user_session.can_ban_users; + let _is_mod = user_session.can_ban_users; + + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + // We should focus on visibility-related controls instead. + const CAN_SEE_DELETED : bool = false; let query_results = list_datasets_by_username( &server_state.mysql_pool, &user_session.username, - is_mod, + CAN_SEE_DELETED, ).await.map_err(|e| { warn!("Error querying for datasets: {:?}", e); ListDatasetsByUserError::ServerError diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_user.rs index a946d292c1..c520103039 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_user.rs @@ -89,12 +89,16 @@ pub async fn list_datasets_by_user_handler( let username = path.username.as_ref(); let creator_user_token = user_session.user_token.clone(); - let is_mod = user_session.can_ban_users; + let _is_mod = user_session.can_ban_users; + + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + // We should focus on visibility-related controls instead. + const CAN_SEE_DELETED : bool = false; let query_results = list_datasets_by_username( &server_state.mysql_pool, &username, - is_mod, + CAN_SEE_DELETED, ).await.map_err(|e| { warn!("Error querying for datasets: {:?}", e); ListDatasetsByUserError::ServerError diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_available_voices.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_available_voices.rs index e5ebe325d9..1ef494e1c8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_available_voices.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_available_voices.rs @@ -110,13 +110,15 @@ pub async fn list_available_voices_handler( let include_user_hidden = is_mod; + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + const CAN_SEE_DELETED : bool = false; let mut query_builder = ListVoicesQueryBuilder::new() .sort_ascending(sort_ascending) .scope_creator_username(None) .include_user_hidden(include_user_hidden) - .include_user_deleted_results(is_mod) - .include_mod_deleted_results(is_mod) + .include_user_deleted_results(CAN_SEE_DELETED) + .include_mod_deleted_results(CAN_SEE_DELETED) .limit(limit) .cursor_is_reversed(cursor_is_reversed) .offset(cursor); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_session.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_session.rs index 3a4cf16409..9a3e6c75ac 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_session.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_session.rs @@ -76,12 +76,16 @@ pub async fn list_voices_by_session_handler( })?; let creator_user_token = user_session.user_token.clone(); - let is_mod = user_session.can_ban_users; + let _is_mod = user_session.can_ban_users; + + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + // We should focus on visibility-related controls instead. + const CAN_SEE_DELETED : bool = false; let query_results = list_zs_voices_by_username( &server_state.mysql_pool, &user_session.username, - is_mod, + CAN_SEE_DELETED, ).await.map_err(|e| { warn!("Error querying for voices: {:?}", e); ListVoicesByUserError::ServerError diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_user.rs index e76429d208..130b044377 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_user.rs @@ -89,12 +89,16 @@ pub async fn list_voices_by_user_handler( let username = path.username.as_ref(); let creator_user_token = user_session.user_token.clone(); - let is_mod = user_session.can_ban_users; + let _is_mod = user_session.can_ban_users; + + // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying! + // We should focus on visibility-related controls instead. + const CAN_SEE_DELETED : bool = false; let query_results = list_zs_voices_by_username( &server_state.mysql_pool, &username, - is_mod, + CAN_SEE_DELETED, ).await.map_err(|e| { warn!("Error querying for voices: {:?}", e); ListVoicesByUserError::ServerError From 380933d1aabec9fc08ca70682cf1758152d93cce Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 18 Jan 2024 22:35:51 -0500 Subject: [PATCH 5/7] much cleaner queries --- .../src/searches/search_model_weights.rs | 241 +++++++++++++++++- 1 file changed, 231 insertions(+), 10 deletions(-) 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 1a2553e593..a03268434b 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -2,10 +2,95 @@ use elasticsearch::{Elasticsearch, SearchParts}; use serde_json::{json, Value}; use enums::by_table::model_weights::weights_types::WeightsType; -use errors::AnyhowResult; +use errors::{anyhow, AnyhowResult}; use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument}; +pub struct SearchModelWeightsQuery<'a> { + pub search_term: &'a str, + pub maybe_language_subtag: Option<&'a str>, + pub maybe_weights_type: Option, +} + +impl <'a>SearchModelWeightsQuery<'a> { + + 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 + } + } + )); + } + + if let Some(weights_type) = &self.maybe_weights_type { + must_clause.push(json!( + { + "match": { + "weights_type": weights_type.to_string(), + } + } + )); + } + + Ok(search_json) + } +} + pub async fn search_model_weights( client: &Elasticsearch, search_term: &str, @@ -13,12 +98,14 @@ pub async fn search_model_weights( maybe_weights_type: Option, ) -> AnyhowResult> { - let search_json = match (maybe_language_subtag, maybe_weights_type) { - (Some(language), None) => query_model_weights_with_required_language(search_term, language), - (None, Some(weights_type)) => query_model_weights_with_model_weights_type(search_term, weights_type), - _ => query_model_weights(search_term), + let query = SearchModelWeightsQuery { + search_term, + maybe_language_subtag, + maybe_weights_type, }; + let search_json = query.query()?; + let search_response = client .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX])) .body(search_json) @@ -102,6 +189,52 @@ fn query_model_weights(search_term: &str) -> Value { }) } +fn base_search(search_term: &str) -> Value { + json!({ + "query": { + "bool": { + "must": [ + { + "bool": { + "should": [ + { + "fuzzy": { + "title": { + "value": search_term, + "fuzziness": 2 + } + } + }, + { + "match": { + "title": { + "query": search_term, + "boost": 1 + } + } + }, + { + "multi_match": { + "query": search_term, + "type": "bool_prefix", + "fields": [ + "title", + "title._2gram", + "title._3gram" + ], + "boost": 50 + } + } + ] + } + } + ] + } + } + }) +} + + fn query_model_weights_with_required_language(search_term: &str, language_tag: &str) -> Value { json!({ "query": { @@ -205,12 +338,31 @@ fn query_model_weights_with_model_weights_type(search_term: &str, weight_type: W #[cfg(test)] mod tests { use regex::Regex; - use crate::searches::search_model_weights::query_model_weights; + use enums::by_table::model_weights::weights_types::WeightsType; + use errors::AnyhowResult; + use crate::searches::search_model_weights::{query_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 json = query_model_weights("FOO_BAR_BAZ"); - let json = serde_json::to_string(&json).unwrap(); + let query = SearchModelWeightsQuery { + search_term: "FOO_BAR_BAZ", + maybe_language_subtag: None, + maybe_weights_type: None, + }; + + let json = query_to_json_string(query).unwrap(); // NB: Keys in emitted JSON are sorted. let expected_json = r#" @@ -258,9 +410,78 @@ mod tests { } "#; - let regex = Regex::new("\\s+").expect("regex should parse"); + let expected_json = compact_json(&expected_json); - let expected_json = regex.replace_all(&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), + }; + + 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": { + "fuzziness": 2, + "value": "FOO_BAR_BAZ" + } + } + }, + { + "match": { + "title": { + "boost": 1, + "query": "FOO_BAR_BAZ" + } + } + }, + { + "multi_match": { + "boost": 50, + "fields": [ + "title", + "title._2gram", + "title._3gram" + ], + "query": "FOO_BAR_BAZ", + "type": "bool_prefix" + } + } + ] + } + }, + { + "match": { + "maybe_ietf_primary_language_subtag": "en" + } + }, + { + "match": { + "weights_type": "so_vits_svc" + } + } + ] + } + } + } + "#; + + let expected_json = compact_json(&expected_json); assert_eq!(&json, &expected_json); } From c44cc5671f45ffe0fa44d9058d62de2501293db7 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 18 Jan 2024 22:53:17 -0500 Subject: [PATCH 6/7] ES cleanup, additional param --- crates/cli/elasticsearch-cli/src/main.rs | 2 +- .../src/searches/search_model_weights.rs | 217 ++---------------- .../weights/search_model_weights_handler.rs | 3 +- 3 files changed, 28 insertions(+), 194 deletions(-) diff --git a/crates/cli/elasticsearch-cli/src/main.rs b/crates/cli/elasticsearch-cli/src/main.rs index dbe8a179c2..3edd5d8cba 100644 --- a/crates/cli/elasticsearch-cli/src/main.rs +++ b/crates/cli/elasticsearch-cli/src/main.rs @@ -57,7 +57,7 @@ pub async fn main() -> AnyhowResult<()> { Action::SearchModelWeights => { info!("Searching model weights..."); //let results = search_model_weights(&elasticsearch, "zel", Some("en"), Some(WeightsType::SoVitsSvc)).await?; - let results = search_model_weights(&elasticsearch, "zel", None, Some(WeightsType::RvcV2)).await?; + let results = search_model_weights(&elasticsearch, "zel", None, Some(WeightsType::RvcV2), None).await?; for result in results { println!("Result: {:#?}", result); } 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 a03268434b..2c3a661238 100644 --- a/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs +++ b/crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs @@ -1,6 +1,7 @@ use elasticsearch::{Elasticsearch, SearchParts}; use serde_json::{json, Value}; +use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use errors::{anyhow, AnyhowResult}; @@ -10,6 +11,7 @@ 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, } impl <'a>SearchModelWeightsQuery<'a> { @@ -87,6 +89,16 @@ impl <'a>SearchModelWeightsQuery<'a> { )); } + if let Some(weights_category) = &self.maybe_weights_category { + must_clause.push(json!( + { + "match": { + "weights_category": weights_category.to_string(), + } + } + )); + } + Ok(search_json) } } @@ -96,12 +108,14 @@ pub async fn search_model_weights( search_term: &str, maybe_language_subtag: Option<&str>, maybe_weights_type: Option, + maybe_weights_category: Option, ) -> AnyhowResult> { let query = SearchModelWeightsQuery { search_term, maybe_language_subtag, maybe_weights_type, + maybe_weights_category, }; let search_json = query.query()?; @@ -144,203 +158,15 @@ pub async fn search_model_weights( Ok(documents) } -fn query_model_weights(search_term: &str) -> Value { - json!({ - "query": { - "bool": { - "must": [ - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": search_term, - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": search_term, - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - }) -} - -fn base_search(search_term: &str) -> Value { - json!({ - "query": { - "bool": { - "must": [ - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": search_term, - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": search_term, - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - }) -} - - -fn query_model_weights_with_required_language(search_term: &str, language_tag: &str) -> Value { - json!({ - "query": { - "bool": { - "must": [ - { - "match": { - "maybe_ietf_primary_language_subtag": language_tag - } - }, - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": search_term, - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": search_term, - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - }) -} - -fn query_model_weights_with_model_weights_type(search_term: &str, weight_type: WeightsType) -> Value { - json!({ - "query": { - "bool": { - "must": [ - { - "match": { - "weights_type": weight_type.to_string(), - } - }, - { - "bool": { - "should": [ - { - "fuzzy": { - "title": { - "value": search_term, - "fuzziness": 2 - } - } - }, - { - "match": { - "title": { - "query": search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": search_term, - "type": "bool_prefix", - "fields": [ - "title", - "title._2gram", - "title._3gram" - ], - "boost": 50 - } - } - ] - } - } - ] - } - } - }) -} - #[cfg(test)] mod 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::{query_model_weights, SearchModelWeightsQuery}; + + use crate::searches::search_model_weights::SearchModelWeightsQuery; fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult { let json = query.query()?; @@ -360,6 +186,7 @@ mod tests { 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(); @@ -421,6 +248,7 @@ mod tests { 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(); @@ -474,6 +302,11 @@ mod tests { "match": { "weights_type": "so_vits_svc" } + }, + { + "match": { + "weights_category": "vocoder" + } } ] } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index aab82b79bd..9788f8b03d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -117,7 +117,8 @@ pub async fn search_model_weights_handler( &server_state.elasticsearch, &request.search_term, request.ietf_language_subtag.as_deref(), - request.weight_type) + request.weight_type, + request.weight_category) .await .map_err(|err| { error!("Searching error: {:?}", err); From a4b549af19025389419304b646f3f2fc6e4eebd2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 19 Jan 2024 02:56:46 -0500 Subject: [PATCH 7/7] allow larger file uploads --- .../src/http_server/web_utils/read_multipart_field_bytes.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/service/web/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs b/crates/service/web/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs index 57fa2fd81f..d9a2325a26 100644 --- a/crates/service/web/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs +++ b/crates/service/web/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs @@ -6,7 +6,8 @@ use futures::{StreamExt, TryStreamExt}; use log::warn; const MIN_BYTES : usize = 10; -const MAX_BYTES : usize = 1024 * 1024 * 20; +//const MAX_BYTES : usize = 1024 * 1024 * 20; // 20 Mb +const MAX_BYTES : usize = 1024 * 1024 * 125; // 125 Mb /// Read any field from a multipart form as bytes, but do some error checking. pub async fn checked_read_multipart_bytes(field: &mut Field) -> anyhow::Result> {