diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs index e80f975507..b81b67e1b5 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs @@ -17,6 +17,7 @@ use enums::common::job_status_plus::JobStatusPlus; use mysql_queries::queries::generic_inference::web::get_inference_job_status::get_inference_job_status; use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use crate::http_server::responses::filter_model_name::maybe_filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; @@ -203,7 +204,7 @@ pub async fn get_inference_job_status_handler( job_token: record.job_token, request: RequestDetailsResponse { inference_category: record.request_details.inference_category, - maybe_model_type: filter_model_name(record.request_details.maybe_model_type), + maybe_model_type: maybe_filter_model_name(record.request_details.maybe_model_type.as_deref()), maybe_model_token: record.request_details.maybe_model_token, maybe_model_title: record.request_details.maybe_model_title, maybe_raw_inference_text: record.request_details.maybe_raw_inference_text, @@ -211,7 +212,7 @@ pub async fn get_inference_job_status_handler( status: StatusDetailsResponse { status: record.status, maybe_extra_status_description, - maybe_assigned_worker: filter_model_name(record.maybe_assigned_worker), + maybe_assigned_worker: maybe_filter_model_name(record.maybe_assigned_worker.as_deref()), maybe_assigned_cluster: record.maybe_assigned_cluster, maybe_first_started_at: record.maybe_first_started_at, attempt_count: record.attempt_count as u8, @@ -295,26 +296,3 @@ pub async fn get_inference_job_status_handler( .content_type("application/json") .body(body)) } - -fn filter_model_name(name: Option) -> Option { - // We're not revealing some of the models we use - name.map(|name| { - name.replace("sadtalker", "faceanimator") - .replace("sad-talker", "face-animator") - .replace("sad_talker", "face_animator") - .replace("vall_e_x", "voice_designer") - }) -} - -#[cfg(test)] -mod tests { - use crate::http_server::endpoints::inference_job::get_inference_job_status::filter_model_name; - - #[test] - fn test_filter_model_name() { - assert_eq!(filter_model_name(None), None); - assert_eq!(filter_model_name(Some("foobarbaz".to_string())), Some("foobarbaz".to_string())); - assert_eq!(filter_model_name(Some("inference-job-sadtalker-5df55cfbb7-ngxzh".to_string())), Some("inference-job-faceanimator-5df55cfbb7-ngxzh".to_string())); - assert_eq!(filter_model_name(Some("sad_talker".to_string())), Some("face_animator".to_string())); - } -} \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs index ce715e6263..c1055c026f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs @@ -14,6 +14,7 @@ use mysql_queries::queries::generic_inference::web::get_inference_job_status::ge use mysql_queries::queries::tts::tts_inference_jobs::get_tts_inference_job_status::get_tts_inference_job_status; use redis_common::redis_keys::RedisKeys; use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use crate::http_server::responses::filter_model_name::filter_model_name; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; @@ -225,6 +226,10 @@ async fn modern_lookup(job_token: &str, server_state: &ServerState) } }; + // NB: Model type is probably TT2, but let's filter it in case a hidden model type ever sneaks in + let model_type = record.request_details.maybe_model_type.as_deref().unwrap_or_else(|| "tacotron2"); + let model_type = filter_model_name(model_type); + Ok(TtsInferenceJobStatusForResponse { job_token: record.job_token.to_string(), status: record.status.to_string(), @@ -251,7 +256,7 @@ async fn modern_lookup(job_token: &str, server_state: &ServerState) } }), model_token: record.request_details.maybe_model_token.unwrap_or_else(|| "NO_MODEL_TOKEN".to_string()), - tts_model_type: record.request_details.maybe_model_type.unwrap_or_else(|| "tacotron2".to_string()), // NB: It's probably TT2 + tts_model_type: model_type, title: record.request_details.maybe_model_title.unwrap_or_else(|| "no model title".to_string()), raw_inference_text: record.request_details.maybe_raw_inference_text.unwrap_or_else(|| "no inference text".to_string()), created_at: record.created_at, diff --git a/crates/service/web/storyteller_web/src/http_server/mod.rs b/crates/service/web/storyteller_web/src/http_server/mod.rs index e3fb34c000..a9ab6c18b9 100644 --- a/crates/service/web/storyteller_web/src/http_server/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/mod.rs @@ -1,4 +1,5 @@ pub mod common_responses; pub mod endpoints; pub mod middleware; +pub mod responses; pub mod web_utils; diff --git a/crates/service/web/storyteller_web/src/http_server/responses/filter_model_name.rs b/crates/service/web/storyteller_web/src/http_server/responses/filter_model_name.rs new file mode 100644 index 0000000000..fea6d59204 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/responses/filter_model_name.rs @@ -0,0 +1,32 @@ + +pub fn filter_model_name(name: &str) -> String { + // We're not revealing some of the models we use + name.replace("sadtalker", "faceanimator") + .replace("sad-talker", "face-animator") + .replace("sad_talker", "face_animator") + .replace("vall-e-x", "voice_designer") + .replace("vall_e_x", "voice_designer") + .replace("vallex", "voice_designer") +} + +pub fn maybe_filter_model_name(name: Option<&str>) -> Option { + name.map(|name| filter_model_name(name)) +} + +#[cfg(test)] +mod tests { + use crate::http_server::responses::filter_model_name::maybe_filter_model_name; + + #[test] + fn test_maybe_filter_model_name() { + // Unfiltered + assert_eq!(maybe_filter_model_name(None), None); + assert_eq!(maybe_filter_model_name(Some("foobarbaz")), Some("foobarbaz".to_string())); + + // Filter raw field values + assert_eq!(maybe_filter_model_name(Some("sad_talker")), Some("face_animator".to_string())); + + // Filter inclusion in pod worker names + assert_eq!(maybe_filter_model_name(Some("inference-job-sadtalker-5df55cfbb7-ngxzh")), Some("inference-job-faceanimator-5df55cfbb7-ngxzh".to_string())); + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/responses/mod.rs b/crates/service/web/storyteller_web/src/http_server/responses/mod.rs new file mode 100644 index 0000000000..8bfd234944 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/responses/mod.rs @@ -0,0 +1 @@ +pub mod filter_model_name; \ No newline at end of file