diff --git a/Cargo.lock b/Cargo.lock index 61f301f0b4..d31846707b 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/cli/elasticsearch-cli/src/main.rs b/crates/cli/elasticsearch-cli/src/main.rs index 06afd4f47e..3edd5d8cba 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), None).await?; for result in results { - println!("Result: {:?}", result); + println!("Result: {:#?}", result); } } } 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 d6caad142f..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,22 +1,125 @@ use elasticsearch::{Elasticsearch, SearchParts}; -use log::error; use serde_json::{json, Value}; -use errors::AnyhowResult; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +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, + pub maybe_weights_category: 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(), + } + } + )); + } + + if let Some(weights_category) = &self.maybe_weights_category { + must_clause.push(json!( + { + "match": { + "weights_category": weights_category.to_string(), + } + } + )); + } + + Ok(search_json) + } +} + pub async fn search_model_weights( client: &Elasticsearch, search_term: &str, maybe_language_subtag: Option<&str>, + maybe_weights_type: Option, + maybe_weights_category: 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 query = SearchModelWeightsQuery { + search_term, + maybe_language_subtag, + maybe_weights_type, + maybe_weights_category, }; + let search_json = query.query()?; + let search_response = client .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX])) .body(search_json) @@ -44,11 +147,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); } } @@ -59,97 +158,164 @@ 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 - } - } - ] - } - } - ] - } - } - }) -} +#[cfg(test)] +mod tests { + use regex::Regex; -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 + 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": { + "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": { - "title": { - "query": search_term, - "boost": 1 - } - } - }, - { - "multi_match": { - "query": search_term, - "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": { + "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" + } + }, + { + "match": { + "weights_category": "vocoder" + } + } + ] + } + } + } + "#; + + let expected_json = compact_json(&expected_json); + + assert_eq!(&json, &expected_json); + } } 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 962d5a93c0..755a952756 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 @@ -247,7 +247,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"); 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 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..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 @@ -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,9 @@ 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, + request.weight_category) .await .map_err(|err| { error!("Searching error: {:?}", err); 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> {