Compare commits

...
Author SHA1 Message Date
Brandon Thomas d0d17756fd codegen 2026-07-06 16:11:16 -04:00
Brandon Thomas 84bf2d934e executor 2026-07-06 16:10:56 -04:00
Brandon Thomas aba271e413 combined query 2026-07-06 15:58:55 -04:00
Brandon Thomas 4fbe9f54e0 more cleanup (#1734) 2026-07-04 11:53:00 -04:00
Brandon Thomas e850022fe8 further cleanup (#1733) 2026-07-04 11:43:49 -04:00
Brandon Thomas 1a528bba45 remove deprecated endpoints (#1732) 2026-07-04 10:15:59 -04:00
Brandon Thomas b625d5ff79 http header fields on alerts (#1731) 2026-07-03 20:59:52 -04:00
Brandon Thomas d8fd0bfe64 fix request logging (#1730) 2026-07-03 16:07:00 -04:00
84 changed files with 1131 additions and 4698 deletions
@@ -1,12 +0,0 @@
{
"db_name": "MySQL",
"query": "\nUPDATE tts_results\nSET\n mod_deleted_at = CURRENT_TIMESTAMP,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 2
},
"nullable": []
},
"hash": "376e36b1357952599866262c7356ba35f15ccdfb338d8ccdbdc06ccd69331fc3"
}
@@ -1,12 +0,0 @@
{
"db_name": "MySQL",
"query": "\nUPDATE tts_results\nSET\n mod_deleted_at = NULL,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 2
},
"nullable": []
},
"hash": "40a3301e9a3201e74d49b6b339f67291e2ae22d315188401708609ac2d28a68c"
}
@@ -1,47 +0,0 @@
{
"db_name": "MySQL",
"query": "\nSELECT\n job_type as queue_type,\n count(*) as pending_job_count,\n NOW() as present_time\n FROM (\n SELECT\n token,\n job_type\n FROM generic_inference_jobs\n WHERE status IN (\"pending\", \"attempt_failed\")\n AND job_type IS NOT NULL\n UNION\n SELECT\n token,\n job_type\n FROM generic_inference_jobs\n WHERE status IN (\"started\")\n AND created_at > (CURDATE() - INTERVAL 15 MINUTE)\n AND job_type IS NOT NULL\n) as generic_inner\nGROUP BY queue_type\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "queue_type",
"type_info": {
"type": "VarString",
"flags": "BINARY",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 1,
"name": "pending_job_count",
"type_info": {
"type": "LongLong",
"flags": "NOT_NULL | BINARY",
"char_set": 63,
"max_size": 21
}
},
{
"ordinal": 2,
"name": "present_time",
"type_info": {
"type": "Datetime",
"flags": "NOT_NULL | BINARY",
"char_set": 63,
"max_size": 19
}
}
],
"parameters": {
"Right": 0
},
"nullable": [
true,
false,
false
]
},
"hash": "73dc054488b2fa3efc923d3a90914a1450b73e99a377562484a0498ccce1676f"
}
@@ -1,12 +0,0 @@
{
"db_name": "MySQL",
"query": "\nINSERT INTO tts_model_upload_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n title = ?,\n tts_model_type = ?,\n download_url = ?,\n status = \"pending\"\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 8
},
"nullable": []
},
"hash": "886861e91c7e8af53eb856c949389a172b8d9e079f1c2070d115423907c47c50"
}
@@ -0,0 +1,102 @@
{
"db_name": "MySQL",
"query": "\nSELECT\n users.token as `user_token: UserToken`,\n users.username,\n users.display_name,\n users.email_address,\n users.user_role_slug,\n users.is_banned,\n user_roles.can_ban_users,\n api_keys.token as `api_key_token: ApiKeyToken`\nFROM api_keys\nINNER JOIN users\n ON users.token = api_keys.owner_user_token\nLEFT OUTER JOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE api_keys.api_key = ?\n AND api_keys.maybe_deleted_at IS NULL\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\nLIMIT 1\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "user_token: UserToken",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 1,
"name": "username",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 2,
"name": "display_name",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 3,
"name": "email_address",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 1020
}
},
{
"ordinal": 4,
"name": "user_role_slug",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 64
}
},
{
"ordinal": 5,
"name": "is_banned",
"type_info": {
"type": "Tiny",
"flags": "NOT_NULL",
"char_set": 63,
"max_size": 1
}
},
{
"ordinal": 6,
"name": "can_ban_users",
"type_info": {
"type": "Tiny",
"flags": "",
"char_set": 63,
"max_size": 1
}
},
{
"ordinal": 7,
"name": "api_key_token: ApiKeyToken",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
}
],
"parameters": {
"Right": 1
},
"nullable": [
false,
false,
false,
false,
false,
false,
true,
false
]
},
"hash": "9188e235caa439d50628333c2f4722b9f1a8f575c192d4d36a2039709c412ac0"
}
@@ -1,12 +0,0 @@
{
"db_name": "MySQL",
"query": "\nUPDATE tts_results\nSET\n user_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 1
},
"nullable": []
},
"hash": "961f792c24bee6009e3e8a215561db932871a53c48839b8cb80292de00e7b2e1"
}
@@ -1,135 +0,0 @@
{
"db_name": "MySQL",
"query": "\nSELECT\n jobs.token as job_token,\n\n jobs.status,\n jobs.attempt_count,\n jobs.on_success_result_token as maybe_result_token,\n results.public_bucket_wav_audio_path as maybe_public_bucket_wav_audio_path,\n\n jobs.model_token,\n tts.tts_model_type,\n tts.title,\n\n jobs.raw_inference_text,\n\n jobs.created_at,\n jobs.updated_at\n\nFROM tts_inference_jobs as jobs\nJOIN tts_models as tts\n ON tts.token = jobs.model_token\nLEFT OUTER JOIN tts_results as results\n ON jobs.on_success_result_token = results.token\n\nWHERE jobs.token = ?\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "job_token",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 1,
"name": "status",
"type_info": {
"type": "String",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | ENUM",
"char_set": 224,
"max_size": 64
}
},
{
"ordinal": 2,
"name": "attempt_count",
"type_info": {
"type": "Long",
"flags": "NOT_NULL",
"char_set": 63,
"max_size": 3
}
},
{
"ordinal": 3,
"name": "maybe_result_token",
"type_info": {
"type": "VarString",
"flags": "MULTIPLE_KEY | BINARY",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 4,
"name": "maybe_public_bucket_wav_audio_path",
"type_info": {
"type": "VarString",
"flags": "BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 1020
}
},
{
"ordinal": 5,
"name": "model_token",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 6,
"name": "tts_model_type",
"type_info": {
"type": "String",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | ENUM",
"char_set": 224,
"max_size": 60
}
},
{
"ordinal": 7,
"name": "title",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 1020
}
},
{
"ordinal": 8,
"name": "raw_inference_text",
"type_info": {
"type": "Blob",
"flags": "NOT_NULL | BLOB | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 262140
}
},
{
"ordinal": 9,
"name": "created_at",
"type_info": {
"type": "Timestamp",
"flags": "NOT_NULL | BINARY | TIMESTAMP",
"char_set": 63,
"max_size": 19
}
},
{
"ordinal": 10,
"name": "updated_at",
"type_info": {
"type": "Timestamp",
"flags": "NOT_NULL | BINARY | TIMESTAMP | ON_UPDATE_NOW",
"char_set": 63,
"max_size": 19
}
}
],
"parameters": {
"Right": 1
},
"nullable": [
false,
false,
false,
true,
true,
false,
false,
false,
false,
false,
false
]
},
"hash": "a2b586dd2120f29dd26281626f7381e292c821f8e826b9bc584b92ffa45035b2"
}
@@ -1,69 +0,0 @@
{
"db_name": "MySQL",
"query": "\nSELECT\n users.username,\n users.display_name,\n users.email_gravatar_hash as gravatar_hash,\n creator_user_token,\n count(*) as uploaded_count\nFROM tts_models\nJOIN users\nON\n users.token = creator_user_token\nWHERE\n users.is_banned IS FALSE\n AND users.token NOT IN (\"U:DRB1SKEMKHT1J\")\n AND tts_models.user_deleted_at IS NULL\n AND tts_models.mod_deleted_at IS NULL\nGROUP BY creator_user_token\nORDER BY uploaded_count desc\nLIMIT 25;\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "username",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 1,
"name": "display_name",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 2,
"name": "gravatar_hash",
"type_info": {
"type": "String",
"flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 3,
"name": "creator_user_token",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 4,
"name": "uploaded_count",
"type_info": {
"type": "LongLong",
"flags": "NOT_NULL | BINARY",
"char_set": 63,
"max_size": 21
}
}
],
"parameters": {
"Right": 0
},
"nullable": [
false,
false,
false,
false,
false
]
},
"hash": "a82fa4eb36528ebf52af2585cdda203476a4ba3645837c792830474fca2eda4f"
}
@@ -0,0 +1,91 @@
{
"db_name": "MySQL",
"query": "\nSELECT\n users.token as `user_token: UserToken`,\n users.username,\n users.display_name,\n users.email_address,\n users.user_role_slug,\n users.is_banned,\n user_roles.can_ban_users\nFROM user_sessions\nINNER JOIN users\n ON users.token = user_sessions.user_token\nLEFT OUTER JOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE user_sessions.token = ?\n AND user_sessions.deleted_at IS NULL\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\nLIMIT 1\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "user_token: UserToken",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 128
}
},
{
"ordinal": 1,
"name": "username",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 2,
"name": "display_name",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 80
}
},
{
"ordinal": 3,
"name": "email_address",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 1020
}
},
{
"ordinal": 4,
"name": "user_role_slug",
"type_info": {
"type": "VarString",
"flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE",
"char_set": 224,
"max_size": 64
}
},
{
"ordinal": 5,
"name": "is_banned",
"type_info": {
"type": "Tiny",
"flags": "NOT_NULL",
"char_set": 63,
"max_size": 1
}
},
{
"ordinal": 6,
"name": "can_ban_users",
"type_info": {
"type": "Tiny",
"flags": "",
"char_set": 63,
"max_size": 1
}
}
],
"parameters": {
"Right": 1
},
"nullable": [
false,
false,
false,
false,
false,
false,
true
]
},
"hash": "c77126adb36f9e2f0248720eb4d8c1cd9d14719f64a02678d8021e26addd2c18"
}
@@ -1,12 +0,0 @@
{
"db_name": "MySQL",
"query": "\nUPDATE tts_results\nSET\n user_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 1
},
"nullable": []
},
"hash": "dbce51eb1c80dded902c9509e50713483c7131f115e8ec10646e99e88110278b"
}
@@ -82,6 +82,7 @@ fn do_build_cors_config(is_production: bool) -> Cors {
actix_http::header::ACCESS_CONTROL_ALLOW_CREDENTIALS, // https://stackoverflow.com/a/46412839
actix_http::header::HeaderName::from_static("x-requested-with"), // Tabulator Ajax sends
actix_http::header::HeaderName::from_static("session"), // Custom header sent by Three.js Storyteller Studio
actix_http::header::HeaderName::from_static("x-artcraft-version"), // ArtCraft client app release, surfaced in error alerts
])
.max_age(3600)
}
@@ -63,6 +63,12 @@ mod tests {
media_file_token: None,
inference_job_token: None,
third_party_id: None,
http_host: None,
http_origin: None,
http_referer: None,
http_user_agent: None,
artcraft_version: None,
trace_id: None,
request_ip_address: None,
avt_cookie_token: None,
session_token: None,
@@ -228,6 +234,12 @@ mod tests {
media_file_token: None,
inference_job_token: None,
third_party_id: None,
http_host: None,
http_origin: None,
http_referer: None,
http_user_agent: None,
artcraft_version: None,
trace_id: None,
request_ip_address: None,
avt_cookie_token: None,
session_token: None,
@@ -37,6 +37,27 @@ pub struct NotificationDetails {
/// HTTP status code associated with the event, if any.
pub(crate) http_status_code: Option<u16>,
/// Hostname the client used to reach this server (from the `Host` or
/// forwarded-host headers), if any. Useful when a server is reachable
/// via multiple domains (eg. api.storyteller.ai, api.fakeyou.com).
pub(crate) http_host: Option<String>,
/// `Origin` header of the request, if any. For browser XHR/fetch this is
/// the scheme+host of the calling frontend.
pub(crate) http_origin: Option<String>,
/// `Referer` header of the request, if any. For browser XHR/fetch this is
/// the URL of the page that made the call (possibly truncated to the
/// origin by the browser's referrer policy).
pub(crate) http_referer: Option<String>,
/// `User-Agent` header of the request, if any.
pub(crate) http_user_agent: Option<String>,
/// `X-ArtCraft-Version` header of the request, if any. Sent by ArtCraft
/// clients to identify which app release made the call.
pub(crate) artcraft_version: Option<String>,
/// User token associated with the event, if any.
pub(crate) user_token: Option<String>,
@@ -188,6 +209,26 @@ impl NotificationDetails {
parts.push(format!("HTTP Status Code: {}", status_code));
}
if let Some(host) = &self.http_host {
parts.push(format!("HTTP Host: {}", host));
}
if let Some(origin) = &self.http_origin {
parts.push(format!("HTTP Origin: {}", origin));
}
if let Some(referer) = &self.http_referer {
parts.push(format!("HTTP Referer: {}", referer));
}
if let Some(user_agent) = &self.http_user_agent {
parts.push(format!("User Agent: {}", user_agent));
}
if let Some(version) = &self.artcraft_version {
parts.push(format!("ArtCraft Version: {}", version));
}
if let Some(ip) = &self.request_ip_address {
parts.push(format!("Request IP: {}", ip));
}
@@ -21,6 +21,11 @@ pub struct NotificationDetailsBuilder {
pub(crate) http_method: Option<String>,
pub(crate) http_path: Option<String>,
pub(crate) http_status_code: Option<u16>,
pub(crate) http_host: Option<String>,
pub(crate) http_origin: Option<String>,
pub(crate) http_referer: Option<String>,
pub(crate) http_user_agent: Option<String>,
pub(crate) artcraft_version: Option<String>,
pub(crate) user_token: Option<String>,
pub(crate) media_file_token: Option<String>,
@@ -53,6 +58,11 @@ impl NotificationDetailsBuilder {
http_method: None,
http_path: None,
http_status_code: None,
http_host: None,
http_origin: None,
http_referer: None,
http_user_agent: None,
artcraft_version: None,
user_token: None,
media_file_token: None,
inference_job_token: None,
@@ -84,6 +94,11 @@ impl NotificationDetailsBuilder {
http_method: None,
http_path: None,
http_status_code: None,
http_host: None,
http_origin: None,
http_referer: None,
http_user_agent: None,
artcraft_version: None,
user_token: None,
media_file_token: None,
inference_job_token: None,
@@ -138,6 +153,31 @@ impl NotificationDetailsBuilder {
self
}
pub fn set_http_host(mut self, http_host: Option<String>) -> Self {
self.http_host = http_host;
self
}
pub fn set_http_origin(mut self, http_origin: Option<String>) -> Self {
self.http_origin = http_origin;
self
}
pub fn set_http_referer(mut self, http_referer: Option<String>) -> Self {
self.http_referer = http_referer;
self
}
pub fn set_http_user_agent(mut self, http_user_agent: Option<String>) -> Self {
self.http_user_agent = http_user_agent;
self
}
pub fn set_artcraft_version(mut self, artcraft_version: Option<String>) -> Self {
self.artcraft_version = artcraft_version;
self
}
pub fn set_user_token(mut self, user_token: Option<String>) -> Self {
self.user_token = user_token;
self
@@ -202,6 +242,11 @@ impl NotificationDetailsBuilder {
http_method: self.http_method,
http_path: self.http_path,
http_status_code: self.http_status_code,
http_host: self.http_host,
http_origin: self.http_origin,
http_referer: self.http_referer,
http_user_agent: self.http_user_agent,
artcraft_version: self.artcraft_version,
user_token: self.user_token,
media_file_token: self.media_file_token,
inference_job_token: self.inference_job_token,
@@ -1,472 +0,0 @@
use std::collections::HashSet;
use chrono::{DateTime, Utc};
use sqlx::{FromRow, MySql, MySqlPool, QueryBuilder, Row};
use sqlx::mysql::MySqlRow;
use enums::by_table::model_weights::{
weights_category::WeightsCategory,
weights_types::WeightsType,
};
use enums::common::view_as::ViewAs;
use enums::traits::mysql_from_row::MySqlFromRow as _;
use enums::common::visibility::Visibility;
use errors::AnyhowResult;
use tokens::tokens::model_weights::ModelWeightToken;
use tokens::tokens::users::UserToken;
pub struct FeaturedWeightsPage {
pub records: Vec<FeaturedWeight>,
pub sort_ascending: bool,
/// ID of the first record in the result set.
pub first_id: Option<i64>,
/// ID of the last record in the result set.
pub last_id: Option<i64>,
}
#[derive(Serialize)]
pub struct FeaturedWeight {
pub token: ModelWeightToken,
pub weights_type: WeightsType,
pub weights_category: WeightsCategory,
pub title: String,
// TODO(bt,2023-12-24): These aren't really appropriate for a list endpoint.
// Hopefully we don't break the frontend by omitting these.
//pub maybe_description_markdown: String,
//pub maybe_description_rendered_html: String,
pub maybe_ietf_language_tag: Option<String>,
pub maybe_ietf_primary_language_subtag: Option<String>,
pub maybe_creator_user_token: Option<UserToken>,
pub maybe_creator_username: Option<String>,
pub maybe_creator_display_name: Option<String>,
pub maybe_creator_email_gravatar_hash: Option<String>,
pub creator_ip_address: String,
pub creator_set_visibility: Visibility,
pub maybe_last_update_user_token: Option<UserToken>,
pub original_download_url: Option<String>,
pub original_filename: Option<String>,
pub file_size_bytes: i32,
pub file_checksum_sha2: String,
pub public_bucket_hash: String,
pub maybe_public_bucket_prefix: Option<String>,
pub maybe_public_bucket_extension: Option<String>,
pub maybe_cover_image_public_bucket_hash: Option<String>,
pub maybe_cover_image_public_bucket_prefix: Option<String>,
pub maybe_cover_image_public_bucket_extension: Option<String>,
pub maybe_ratings_positive_count: Option<u32>,
pub maybe_ratings_negative_count: Option<u32>,
pub maybe_bookmark_count: Option<u32>,
pub cached_usage_count: u64,
pub version: i32,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub user_deleted_at: Option<DateTime<Utc>>,
pub mod_deleted_at: Option<DateTime<Utc>>,
}
pub struct ListFeaturedWeightsArgs<'a> {
pub limit: usize,
pub maybe_offset: Option<usize>,
pub cursor_is_reversed: bool,
pub sort_ascending: bool,
pub view_as: ViewAs,
pub maybe_scoped_weight_types: Option<&'a HashSet<WeightsType>>,
pub maybe_scoped_weight_categories: Option<&'a HashSet<WeightsCategory>>,
pub mysql_pool: &'a MySqlPool,
}
pub async fn list_featured_weights(args: ListFeaturedWeightsArgs<'_>) -> AnyhowResult<FeaturedWeightsPage> {
// TODO(bt,2024-05-05): Why were we doing this, again?
// let count_fields = select_total_count_field();
// let mut count_query_builder = query_builder(
// args.limit,
// args.maybe_offset,
// args.cursor_is_reversed,
// args.sort_ascending,
// count_fields.as_str(),
// args.view_as,
// args.maybe_scoped_weight_type,
// args.maybe_scoped_weight_category,
// );
// let row_count_query = count_query_builder.build_query_scalar::<i64>();
// let row_count_result = row_count_query.fetch_one(args.mysql_pool).await?;
/// Now fetch the actual results with all the fields
let result_fields = select_result_fields();
let mut query = query_builder(
args.limit,
args.maybe_offset,
args.cursor_is_reversed,
args.sort_ascending,
result_fields.as_str(),
args.view_as,
args.maybe_scoped_weight_types,
args.maybe_scoped_weight_categories,
);
let query = query.build_query_as::<RawWeightJoinUser>();
let results = query.fetch_all(args.mysql_pool).await?;
let first_id = results.first()
.map(|raw_result| raw_result.id);
let last_id = results.last()
.map(|raw_result| raw_result.id);
let weights_records: Vec<FeaturedWeight> = map_to_weights(results).await;
Ok(FeaturedWeightsPage {
records: weights_records,
sort_ascending: args.sort_ascending,
first_id,
last_id,
})
}
fn select_total_count_field() -> String {
r#"
COUNT(mw.id) AS total_count
"#
.to_string()
}
fn select_result_fields() -> String {
r#"
mw.id,
mw.token,
mw.title,
mw.weights_type,
mw.weights_category,
mw.maybe_ietf_language_tag,
mw.maybe_ietf_primary_language_subtag,
u.token as maybe_creator_user_token,
u.username as maybe_creator_username,
u.display_name as maybe_creator_display_name,
u.email_gravatar_hash as maybe_creator_email_gravatar_hash,
mw.creator_ip_address,
mw.creator_set_visibility,
mw.maybe_last_update_user_token,
mw.original_download_url,
mw.original_filename,
mw.file_size_bytes,
mw.file_checksum_sha2,
mw.public_bucket_hash,
mw.maybe_public_bucket_prefix,
mw.maybe_public_bucket_extension,
cover_image.public_bucket_directory_hash as maybe_cover_image_public_bucket_hash,
cover_image.maybe_public_bucket_prefix as maybe_cover_image_public_bucket_prefix,
cover_image.maybe_public_bucket_extension as maybe_cover_image_public_bucket_extension,
entity_stats.ratings_positive_count as maybe_ratings_positive_count,
entity_stats.ratings_negative_count as maybe_ratings_negative_count,
entity_stats.bookmark_count as maybe_bookmark_count,
mw.cached_usage_count,
mw.version,
mw.created_at,
mw.updated_at,
mw.user_deleted_at,
mw.mod_deleted_at
"#.to_string()
}
fn query_builder<'a>(
_limit: usize, // TODO(bt,2024-05-16): This seems wrong. Why is this unused?
maybe_offset: Option<usize>,
cursor_is_reversed: bool,
sort_ascending: bool,
select_fields: &'a str,
view_as: ViewAs,
maybe_scoped_weight_types: Option<&'a HashSet<WeightsType>>,
maybe_scoped_weight_categories: Option<&'a HashSet<WeightsCategory>>,
) -> QueryBuilder<'a, MySql> {
// NB: Query cannot be statically checked by sqlx
let mut query_builder: QueryBuilder<MySql> = QueryBuilder::new(
format!(r#"
SELECT
{select_fields}
FROM model_weights as mw
JOIN featured_items
ON featured_items.entity_type = "model_weight"
AND featured_items.deleted_at IS NULL
AND featured_items.entity_token = mw.token
LEFT OUTER JOIN users as u
ON u.token = mw.creator_user_token
LEFT OUTER JOIN media_files as cover_image
ON cover_image.token = mw.maybe_cover_image_media_file_token
LEFT OUTER JOIN entity_stats
ON entity_stats.entity_type = "model_weights"
AND entity_stats.entity_token = mw.token
"#
));
let mut first_predicate_added = false;
match view_as {
ViewAs::Author | ViewAs::Moderator => {
if !first_predicate_added {
query_builder.push(" WHERE ");
first_predicate_added = true;
} else {
query_builder.push(" AND ");
}
// NB(bt): Actually, mods don't want to see deleted models. We'll improve the moderator UI later.
query_builder.push(" mw.user_deleted_at IS NULL AND mw.mod_deleted_at IS NULL ");
}
ViewAs::AnotherUser => {
if !first_predicate_added {
query_builder.push(" WHERE ");
first_predicate_added = true;
} else {
query_builder.push(" AND ");
}
query_builder.push(" mw.user_deleted_at IS NULL AND mw.mod_deleted_at IS NULL ");
query_builder.push(" AND mw.creator_set_visibility = ");
query_builder.push_bind(Visibility::Public.to_str());
}
}
if let Some(weight_types) = maybe_scoped_weight_types {
// NB: `WHERE IN` comma separated syntax will be wrong if list has zero length
// We'll skip the predicate if the list isn't empty.
if !weight_types.is_empty() {
if !first_predicate_added {
query_builder.push(" WHERE ");
first_predicate_added = true;
} else {
query_builder.push(" AND ");
}
query_builder.push(" mw.weights_type IN ( ");
let mut separated = query_builder.separated(", ");
for weight_type in weight_types.iter() {
separated.push_bind(weight_type.to_str());
}
separated.push_unseparated(") ");
}
}
if let Some(weight_categories) = maybe_scoped_weight_categories {
// NB: `WHERE IN` comma separated syntax will be wrong if list has zero length
// We'll skip the predicate if the list isn't empty.
if !weight_categories.is_empty() {
if !first_predicate_added {
query_builder.push(" WHERE ");
first_predicate_added = true;
} else {
query_builder.push(" AND ");
}
query_builder.push(" mw.weights_category IN ( ");
let mut separated = query_builder.separated(", ");
for weights_category in weight_categories.iter() {
separated.push_bind(weights_category.to_str());
}
separated.push_unseparated(") ");
}
}
if let Some(offset) = maybe_offset {
if !first_predicate_added {
query_builder.push(" WHERE ");
first_predicate_added = true;
} else {
query_builder.push(" AND ");
}
let mut sort_ascending = sort_ascending;
if sort_ascending {
if cursor_is_reversed {
// NB: We're searching backwards.
query_builder.push(" m.id < ");
sort_ascending = !sort_ascending;
} else {
query_builder.push(" m.id > ");
}
} else {
if cursor_is_reversed {
// NB: We're searching backwards.
query_builder.push(" m.id > ");
sort_ascending = !sort_ascending;
} else {
query_builder.push(" m.id < ");
}
}
query_builder.push_bind(offset as i64);
}
query_builder
}
async fn map_to_weights(dataset:Vec<RawWeightJoinUser>) -> Vec<FeaturedWeight> {
let weights: Vec<FeaturedWeight> = dataset
.into_iter()
.map(|weight: RawWeightJoinUser| {
FeaturedWeight {
token: weight.token,
title: weight.title,
weights_type: weight.weights_type,
weights_category: weight.weights_category,
maybe_ietf_language_tag: weight.maybe_ietf_language_tag,
maybe_ietf_primary_language_subtag: weight.maybe_ietf_primary_language_subtag,
maybe_creator_user_token: weight.maybe_creator_user_token,
maybe_creator_username: weight.maybe_creator_username,
maybe_creator_display_name: weight.maybe_creator_display_name,
maybe_creator_email_gravatar_hash: weight.maybe_creator_email_gravatar_hash,
creator_ip_address: weight.creator_ip_address,
creator_set_visibility: weight.creator_set_visibility,
maybe_last_update_user_token: weight.maybe_last_update_user_token,
original_download_url: weight.original_download_url,
original_filename: weight.original_filename,
file_size_bytes: weight.file_size_bytes,
file_checksum_sha2: weight.file_checksum_sha2,
public_bucket_hash: weight.public_bucket_hash,
maybe_public_bucket_prefix: weight.maybe_public_bucket_prefix,
maybe_public_bucket_extension: weight.maybe_public_bucket_extension,
maybe_cover_image_public_bucket_hash: weight.maybe_cover_image_public_bucket_hash,
maybe_cover_image_public_bucket_prefix: weight.maybe_cover_image_public_bucket_prefix,
maybe_cover_image_public_bucket_extension: weight.maybe_cover_image_public_bucket_extension,
maybe_ratings_positive_count: weight.maybe_ratings_positive_count,
maybe_ratings_negative_count: weight.maybe_ratings_negative_count,
maybe_bookmark_count: weight.maybe_bookmark_count,
cached_usage_count: weight.cached_usage_count,
version: weight.version,
created_at: weight.created_at,
updated_at: weight.updated_at,
user_deleted_at: weight.user_deleted_at,
mod_deleted_at: weight.mod_deleted_at,
}
}).collect();
weights
}
struct RawWeightJoinUser {
id: i64,
token: ModelWeightToken,
weights_type: WeightsType,
weights_category: WeightsCategory,
title: String,
maybe_ietf_language_tag: Option<String>,
maybe_ietf_primary_language_subtag: Option<String>,
// TODO(bt,2023-12-24): These aren't really appropriate for a list endpoint.
// Hopefully we don't break the frontend by omitting these.
//pub description_markdown: String,
//pub description_rendered_html: String,
maybe_creator_user_token: Option<UserToken>,
maybe_creator_username: Option<String>,
maybe_creator_display_name: Option<String>,
maybe_creator_email_gravatar_hash: Option<String>,
creator_ip_address: String,
creator_set_visibility: Visibility,
maybe_last_update_user_token: Option<UserToken>,
original_download_url: Option<String>,
original_filename: Option<String>,
file_size_bytes: i32,
file_checksum_sha2: String,
public_bucket_hash: String,
maybe_public_bucket_prefix: Option<String>,
maybe_public_bucket_extension: Option<String>,
cached_usage_count: u64,
maybe_cover_image_public_bucket_hash: Option<String>,
maybe_cover_image_public_bucket_prefix: Option<String>,
maybe_cover_image_public_bucket_extension: Option<String>,
maybe_ratings_positive_count: Option<u32>,
maybe_ratings_negative_count: Option<u32>,
maybe_bookmark_count: Option<u32>,
version: i32,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
user_deleted_at: Option<DateTime<Utc>>,
mod_deleted_at: Option<DateTime<Utc>>,
}
impl FromRow<'_, MySqlRow> for RawWeightJoinUser {
fn from_row(row: &MySqlRow) -> Result<Self, sqlx::Error> {
Ok(Self {
id: row.try_get("id")?,
token: row.try_get("token")?,
weights_type: row.try_get("weights_type")?,
weights_category: row.try_get("weights_category")?,
title: row.try_get("title")?,
maybe_ietf_language_tag: row.try_get("maybe_ietf_language_tag")?,
maybe_ietf_primary_language_subtag: row.try_get("maybe_ietf_primary_language_subtag")?,
// TODO(bt,2023-12-24): These aren't really appropriate for a list endpoint.
// Hopefully we don't break the frontend by omitting these.
//maybe_description_markdown: row.try_get("maybe_description_markdown")?,
//maybe_description_rendered_html: row.try_get("maybe_description_rendered_html")?,
maybe_creator_user_token: row.try_get("maybe_creator_user_token")?,
maybe_creator_username: row.try_get("maybe_creator_username")?,
maybe_creator_display_name: row.try_get("maybe_creator_display_name")?,
maybe_creator_email_gravatar_hash: row.try_get("maybe_creator_email_gravatar_hash")?,
creator_ip_address: row.try_get("creator_ip_address")?,
creator_set_visibility: Visibility::try_from_mysql_row(row, "creator_set_visibility")?,
maybe_last_update_user_token: row.try_get("maybe_last_update_user_token")?,
original_download_url: row.try_get("original_download_url")?,
original_filename: row.try_get("original_filename")?,
file_size_bytes: row.try_get("file_size_bytes")?,
file_checksum_sha2: row.try_get("file_checksum_sha2")?,
public_bucket_hash: row.try_get("public_bucket_hash")?,
maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?,
maybe_public_bucket_extension: row.try_get("maybe_public_bucket_extension")?,
cached_usage_count: row.try_get("cached_usage_count")?,
maybe_cover_image_public_bucket_hash: row.try_get("maybe_cover_image_public_bucket_hash")?,
maybe_cover_image_public_bucket_prefix: row.try_get("maybe_cover_image_public_bucket_prefix")?,
maybe_cover_image_public_bucket_extension: row.try_get("maybe_cover_image_public_bucket_extension")?,
maybe_ratings_positive_count: row.try_get("maybe_ratings_positive_count")?,
maybe_ratings_negative_count: row.try_get("maybe_ratings_negative_count")?,
maybe_bookmark_count: row.try_get("maybe_bookmark_count")?,
version: row.try_get("version")?,
created_at: row.try_get("created_at")?,
updated_at: row.try_get("updated_at")?,
user_deleted_at: row.try_get("user_deleted_at")?,
mod_deleted_at: row.try_get("mod_deleted_at")?,
})
}
}
@@ -1,6 +1,5 @@
pub mod list_all_model_weight_tokens_for_backfill;
pub mod list_all_model_weight_tokens_for_cached_usage_backfill;
pub mod list_featured_weights;
pub mod list_model_weight_tokens_updated_since;
pub mod list_model_weights_for_elastic_search_backfill_using_cursor;
pub mod list_model_weights_for_text_to_speech;
@@ -1,66 +0,0 @@
use chrono::NaiveDateTime;
use sqlx::MySqlPool;
use errors::AnyhowResult;
use crate::helpers::numeric_converters::try_i64_to_u64_or_min;
#[derive(Serialize, Deserialize, Clone)]
pub struct QueueStatsRow {
pub queue_type: String,
pub pending_job_count: u64,
pub present_time: NaiveDateTime,
}
struct QueueStatsRowInternal {
queue_type: Option<String>, // FIXME: Shouldn't be nullable
pending_job_count: i64,
present_time: NaiveDateTime,
}
// NB: Returning Option<T> is a hack to fit with outer Redis caching.
// I'm only doing this because I'm in a hurry. Don't do this.
pub async fn get_unified_queue_stats(
pool: &MySqlPool,
) -> AnyhowResult<Vec<QueueStatsRow>> {
// NB (1): We query as a union since "started" jobs can get picked off mid-run and go stale
// forever. We currently have no automated means of collecting those jobs.
// NB (2): We query the server timestamp so we can cache the results.
// The frontend can then monotonically adjust the count based on timestamp by ignoring
// past times.
let rows : Vec<QueueStatsRowInternal> = sqlx::query_as!(
QueueStatsRowInternal,
r#"
SELECT
job_type as queue_type,
count(*) as pending_job_count,
NOW() as present_time
FROM (
SELECT
token,
job_type
FROM generic_inference_jobs
WHERE status IN ("pending", "attempt_failed")
AND job_type IS NOT NULL
UNION
SELECT
token,
job_type
FROM generic_inference_jobs
WHERE status IN ("started")
AND created_at > (CURDATE() - INTERVAL 15 MINUTE)
AND job_type IS NOT NULL
) as generic_inner
GROUP BY queue_type
"#,
)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(|row|
QueueStatsRow {
queue_type: row.queue_type.unwrap_or_else(|| "unknown".to_string()),
pending_job_count: try_i64_to_u64_or_min(row.pending_job_count),
present_time: row.present_time,
}).collect::<Vec<_>>())
}
@@ -1,3 +1,2 @@
pub mod get_on_prem_worker_stats;
pub mod get_unified_queue_stats;
pub mod get_voice_count_stats;
pub mod get_voice_count_stats;
@@ -1,7 +1,6 @@
pub mod stats;
pub mod tts_category_assignments;
pub mod tts_download_jobs;
pub mod tts_inference_jobs;
pub mod tts_model_upload_jobs;
pub mod tts_models;
pub mod tts_results;
pub mod tts_results;
@@ -1,89 +0,0 @@
use anyhow::anyhow;
use log::warn;
use sqlx::MySql;
use sqlx::pool::PoolConnection;
use errors::AnyhowResult;
#[derive(Serialize, Clone)]
pub struct TtsLeaderboardRecordForList {
pub username: String,
pub display_name: String,
pub gravatar_hash: String,
pub creator_user_token: String,
pub uploaded_count: i64,
}
#[derive(Serialize)]
struct TtsLeaderboardRecordForListRaw {
pub username: String,
pub display_name: String,
pub gravatar_hash: String,
pub creator_user_token: String,
pub uploaded_count: i64,
}
// NB: This may need to become an offline query if it is expensive.
pub async fn calculate_tts_model_leaderboard(
mysql_pool: &mut PoolConnection<MySql>
) -> AnyhowResult<Vec<TtsLeaderboardRecordForList>> {
// NB(1): We're not requiring "is_public_listing_approved IS TRUE" for TTS models
// This is the opposite for W2L templates (for now at least)
// NB(2): We're excluding a few users (eg. Vegito's alternate account)
// U:DRB1SKEMKHT1J - https://fakeyou.com/profile/bluefusion/uploaded-tts-models
let maybe_results = sqlx::query_as!(
TtsLeaderboardRecordForListRaw,
r#"
SELECT
users.username,
users.display_name,
users.email_gravatar_hash as gravatar_hash,
creator_user_token,
count(*) as uploaded_count
FROM tts_models
JOIN users
ON
users.token = creator_user_token
WHERE
users.is_banned IS FALSE
AND users.token NOT IN ("U:DRB1SKEMKHT1J")
AND tts_models.user_deleted_at IS NULL
AND tts_models.mod_deleted_at IS NULL
GROUP BY creator_user_token
ORDER BY uploaded_count desc
LIMIT 25;
"#,
)
.fetch_all(&mut **mysql_pool)
.await;
let results : Vec<TtsLeaderboardRecordForListRaw> = match maybe_results {
Ok(results) => results,
Err(err) => {
warn!("Error: {:?}", err);
match err {
sqlx::Error::RowNotFound => {
return Ok(Vec::new());
},
_ => {
warn!("tts leaderboard query error: {:?}", err);
return Err(anyhow!("tts leaderboard query error"));
}
}
}
};
Ok(results.into_iter()
.map(|result| {
TtsLeaderboardRecordForList {
username: result.username.to_string(),
display_name: result.display_name.to_string(),
gravatar_hash: result.gravatar_hash.to_string(),
creator_user_token: result.creator_user_token.to_string(),
uploaded_count: result.uploaded_count,
}
})
.collect::<Vec<TtsLeaderboardRecordForList>>())
}
@@ -1 +0,0 @@
pub mod calculate_tts_model_leaderboard;
@@ -1,279 +0,0 @@
use chrono::{DateTime, Utc};
use sqlx::{MySql, MySqlPool};
use sqlx::pool::PoolConnection;
use errors::AnyhowResult;
use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_optional_bool};
/// List of assigned categories for a TTS model
/// This is *NOT* paginated, as that defeats the purpose of having a category system.
#[derive(Serialize)]
pub struct AssignedCategoryList {
pub categories: Vec<AssignedCategory>,
}
#[derive(Serialize)]
pub struct AssignedCategory {
pub category_token: String,
pub model_type: String, // TODO: ENUM
pub maybe_super_category_token: Option<String>,
pub can_directly_have_models: bool,
pub can_have_subcategories: bool,
pub can_only_mods_apply: bool,
pub name: String,
pub maybe_dropdown_name: Option<String>,
pub category_creator_user_token: Option<String>,
pub category_creator_username: Option<String>,
pub category_creator_display_name: Option<String>,
pub category_creator_gravatar_hash: Option<String>,
// Moderator fields
pub is_mod_approved: Option<bool>,
pub maybe_mod_comments: Option<String>,
//pub maybe_mod_user_token: Option<String>,
//pub maybe_mod_username: Option<String>,
//pub maybe_mod_display_name: Option<String>,
pub category_created_at: DateTime<Utc>,
pub category_updated_at: DateTime<Utc>,
pub category_deleted_at: Option<DateTime<Utc>>,
pub category_assignment_created_at: DateTime<Utc>,
pub category_assignment_updated_at: DateTime<Utc>,
pub category_assignment_deleted_at: Option<DateTime<Utc>>,
}
/// Query Builder for listing categories
/// These are very difficult queries, so the builder helps for
/// testability, construction, and correctness.
pub struct ListAssignedTtsCategoriesQueryBuilder {
// REQUIRED.
scope_tts_model_token: String,
// Sometimes "can_directly_have_models" categories change state.
// Mods will want to see invalid uses.
show_invalid_model_not_allowed_categories: bool,
// The public lists won't include unapproved, but mod views
// will want to see them.
show_unapproved: bool,
// The public lists won't include deleted models, but mod views
// will want to see them.
show_deleted: bool,
}
impl ListAssignedTtsCategoriesQueryBuilder {
pub fn for_model_token(model_token: &str) -> Self {
Self {
scope_tts_model_token: model_token.to_string(),
show_invalid_model_not_allowed_categories: false,
show_unapproved: false,
show_deleted: false,
}
}
pub fn show_invalid_model_not_allowed_categories(mut self, show_invalid_model_not_allowed_categories: bool) -> Self {
self.show_invalid_model_not_allowed_categories = show_invalid_model_not_allowed_categories;
self
}
pub fn show_unapproved(mut self, show_unapproved: bool) -> Self {
self.show_unapproved = show_unapproved;
self
}
pub fn show_deleted(mut self, show_deleted: bool) -> Self {
self.show_deleted = show_deleted;
self
}
/// Perform the query based on the set predicates.
#[deprecated = "Use the PoolConnection<MySql> method instead of the MySqlPool one."]
pub async fn perform_query(
&self,
mysql_pool: &MySqlPool
) -> AnyhowResult<AssignedCategoryList> {
let mut mysql_connection = mysql_pool.acquire().await?;
self.perform_query_by_connection(&mut mysql_connection).await
}
/// Perform the query based on the set predicates.
pub async fn perform_query_by_connection(
&self,
mysql_connection: &mut PoolConnection<MySql>
) -> AnyhowResult<AssignedCategoryList> {
let internal_results = self.perform_internal_query(mysql_connection).await?;
let categories = internal_results
.into_iter()
.map(|c| {
AssignedCategory {
category_token: c.category_token,
model_type: c.model_type,
maybe_super_category_token: c.maybe_super_category_token,
can_directly_have_models: i8_to_bool(c.can_directly_have_models),
can_have_subcategories: i8_to_bool(c.can_have_subcategories),
can_only_mods_apply: i8_to_bool(c.can_only_mods_apply),
name: c.name,
maybe_dropdown_name: c.maybe_dropdown_name,
category_creator_user_token: c.category_creator_user_token,
category_creator_username: c.category_creator_username,
category_creator_display_name: c.category_creator_display_name,
category_creator_gravatar_hash: c.category_creator_gravatar_hash,
is_mod_approved: nullable_i8_to_optional_bool(c.is_mod_approved),
maybe_mod_comments: c.maybe_mod_comments,
category_created_at: c.category_created_at,
category_updated_at: c.category_updated_at,
category_deleted_at: c.category_deleted_at,
category_assignment_created_at: c.category_assignment_created_at,
category_assignment_updated_at: c.category_assignment_updated_at,
category_assignment_deleted_at: c.category_assignment_deleted_at,
}
})
.collect::<Vec<AssignedCategory>>();
Ok(AssignedCategoryList {
categories,
})
}
/// Perform the query based on the set predicates.
async fn perform_internal_query(
&self,
mysql_connection: &mut PoolConnection<MySql>
) -> AnyhowResult<Vec<RawInternalCategoryRecord>> {
let query = self.build_query_string();
let mut query = sqlx::query_as::<_, RawInternalCategoryRecord>(&query);
// NB: The following bindings must match the order of the query builder !!
// NB: Model token is always used for scoping.
query = query.bind(&self.scope_tts_model_token);
let results = query.fetch_all(&mut **mysql_connection)
.await?;
Ok(results)
}
fn build_query_string(&self) -> String {
// TODO: I haven't figured out how to get field name disambiguation and type coercion working here.
// (1) tts_results.creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`,
// Query error: no column found for name: creator_set_visibility
// (2) creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`,
// Column 'creator_set_visibility' in field list is ambiguous
// TODO/NB: Unfortunately SQLx can't statically typecheck this query
let mut query = r#"
SELECT
model_categories.token as category_token,
model_categories.model_type,
model_categories.maybe_super_category_token,
model_categories.can_directly_have_models,
model_categories.can_have_subcategories,
model_categories.can_only_mods_apply,
model_categories.name,
model_categories.maybe_dropdown_name,
users.token as category_creator_user_token,
users.username as category_creator_username,
users.display_name as category_creator_display_name,
users.email_gravatar_hash as category_creator_gravatar_hash,
model_categories.is_mod_approved,
model_categories.maybe_mod_comments,
model_categories.created_at as category_created_at,
model_categories.updated_at as category_updated_at,
model_categories.deleted_at as category_deleted_at,
tts_category_assignments.created_at as category_assignment_created_at,
tts_category_assignments.updated_at as category_assignment_updated_at,
tts_category_assignments.deleted_at as category_assignment_deleted_at
FROM model_categories
JOIN tts_category_assignments
ON tts_category_assignments.category_token = model_categories.token
LEFT OUTER JOIN users
ON model_categories.creator_user_token = users.token
"#.to_string();
query.push_str(&self.build_predicates());
query
}
pub fn build_predicates(&self) -> String {
let mut query = "".to_string();
// Scoping to tts model is required.
query.push_str(" WHERE tts_category_assignments.model_token = ?");
query.push_str(" AND model_categories.model_type = 'tts'");
// We also don't care to ever show soft-deleted assignments
query.push_str(" AND tts_category_assignments.deleted_at IS NULL");
if !self.show_invalid_model_not_allowed_categories {
query.push_str(" AND model_categories.can_directly_have_models IS TRUE");
}
if !self.show_unapproved {
query.push_str(" AND model_categories.is_mod_approved IS TRUE");
}
if !self.show_deleted {
query.push_str(" AND model_categories.deleted_at IS NULL");
}
query
}
}
#[derive(sqlx::FromRow)]
pub struct RawInternalCategoryRecord {
pub category_token: String,
pub model_type: String, // TODO: ENUM
pub maybe_super_category_token: Option<String>,
pub can_directly_have_models: i8,
pub can_have_subcategories: i8,
pub can_only_mods_apply: i8,
pub name: String,
pub maybe_dropdown_name: Option<String>,
pub category_creator_user_token: Option<String>,
pub category_creator_username: Option<String>,
pub category_creator_display_name: Option<String>,
pub category_creator_gravatar_hash: Option<String>,
// Moderator fields
pub is_mod_approved: Option<i8>,
pub maybe_mod_comments: Option<String>,
pub category_created_at: DateTime<Utc>,
pub category_updated_at: DateTime<Utc>,
pub category_deleted_at: Option<DateTime<Utc>>,
pub category_assignment_created_at: DateTime<Utc>,
pub category_assignment_updated_at: DateTime<Utc>,
pub category_assignment_deleted_at: Option<DateTime<Utc>>,
}
#[cfg(test)]
mod tests {
}
@@ -1,2 +1 @@
pub mod fetch_and_build_tts_model_category_map;
pub mod list_assigned_tts_categories_query_builder;
@@ -1,74 +0,0 @@
use anyhow::anyhow;
use chrono::{DateTime, Utc};
use sqlx::MySqlPool;
use errors::AnyhowResult;
#[derive(Serialize)]
pub struct TtsInferenceJobStatusRecord {
pub job_token: String,
pub status: String,
pub attempt_count: i32,
pub maybe_result_token: Option<String>,
pub maybe_public_bucket_wav_audio_path: Option<String>,
pub model_token: String,
pub tts_model_type: String,
pub title: String, // Name of the TTS model
pub raw_inference_text: String, // User text
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
pub async fn get_tts_inference_job_status(
tts_inference_job_token: &str,
mysql_pool: &MySqlPool
) -> AnyhowResult<Option<TtsInferenceJobStatusRecord>> {
// NB: Lookup failure is Err(RowNotFound).
// NB: Since this is publicly exposed, we don't query sensitive data.
let maybe_status = sqlx::query_as!(
TtsInferenceJobStatusRecord,
r#"
SELECT
jobs.token as job_token,
jobs.status,
jobs.attempt_count,
jobs.on_success_result_token as maybe_result_token,
results.public_bucket_wav_audio_path as maybe_public_bucket_wav_audio_path,
jobs.model_token,
tts.tts_model_type,
tts.title,
jobs.raw_inference_text,
jobs.created_at,
jobs.updated_at
FROM tts_inference_jobs as jobs
JOIN tts_models as tts
ON tts.token = jobs.model_token
LEFT OUTER JOIN tts_results as results
ON jobs.on_success_result_token = results.token
WHERE jobs.token = ?
"#,
tts_inference_job_token
)
.fetch_one(mysql_pool)
.await; // TODO: This will return error if it doesn't exist
match maybe_status {
Ok(record) => Ok(Some(record)),
Err(err) => match err {
sqlx::Error::RowNotFound => Ok(None),
_ => Err(anyhow!("tts job query error: {:?}", err)),
}
}
}
@@ -1,6 +1,5 @@
pub mod _keys;
pub mod get_pending_tts_inference_job_detailed_stats;
pub mod get_tts_inference_job_status;
pub mod insert_tts_inference_job;
pub mod kill_tts_inference_jobs;
pub mod list_available_tts_inference_jobs;
@@ -1,82 +0,0 @@
use anyhow::anyhow;
use log::warn;
use sqlx::error::Error::Database;
use sqlx::MySqlPool;
use errors::AnyhowResult;
use tokens::tokens::tts_model_upload_jobs::TtsModelUploadJobToken;
pub struct InsertTtsModelUploadJobArgs<'a> {
pub uuid: &'a str,
pub creator_user_token: &'a str,
pub creator_ip_address: &'a str,
pub creator_set_visibility: &'a str,
pub title: &'a str,
pub tts_model_type: &'a str,
pub download_url: &'a str,
pub mysql_pool: &'a MySqlPool,
}
// NB: Returns job token
pub async fn insert_tts_model_upload_job(args: InsertTtsModelUploadJobArgs<'_>) -> AnyhowResult<String> {
// This token is returned to the client.
let job_token = TtsModelUploadJobToken::generate().to_string();
let query_result = sqlx::query!(
r#"
INSERT INTO tts_model_upload_jobs
SET
token = ?,
uuid_idempotency_token = ?,
creator_user_token = ?,
creator_ip_address = ?,
creator_set_visibility = ?,
title = ?,
tts_model_type = ?,
download_url = ?,
status = "pending"
"#,
job_token.to_string(),
args.uuid,
args.creator_user_token,
args.creator_ip_address,
args.creator_set_visibility,
args.title,
args.tts_model_type,
args.download_url,
)
.execute(args.mysql_pool)
.await;
let _record_id = match query_result {
Ok(res) => {
res.last_insert_id()
},
Err(err) => {
warn!("New tts model upload creation DB error: {:?}", err);
// NB: SQLSTATE[23000]: Integrity constraint violation
// NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access)
match err {
Database(ref err) => {
let _maybe_code = err.code().map(|c| c.into_owned());
/*match maybe_code.as_deref() {
Some("23000") => {
if err.message().contains("username") {
return Err(UsernameTaken);
} else if err.message().contains("email_address") {
return Err(EmailTaken);
}
}
_ => {},
}*/
},
_ => {},
}
return Err(anyhow!("error inserting model upload job record: {:?}", err));
}
};
Ok(job_token)
}
@@ -1,2 +1 @@
pub mod get_tts_model_upload_job_status;
pub mod insert_tts_model_upload_job;
@@ -1,89 +0,0 @@
use sqlx::MySqlPool;
use errors::AnyhowResult;
pub async fn delete_tts_inference_result_as_user(
inference_result_token: &str,
mysql_pool: &MySqlPool
) -> AnyhowResult<()> {
let _r = sqlx::query!(
r#"
UPDATE tts_results
SET
user_deleted_at = CURRENT_TIMESTAMP
WHERE
token = ?
LIMIT 1
"#,
inference_result_token,
)
.execute(mysql_pool)
.await?;
Ok(())
}
pub async fn delete_tts_inference_result_as_mod(
inference_result_token: &str,
mod_user_token: &str,
mysql_pool: &MySqlPool
) -> AnyhowResult<()> {
let _r = sqlx::query!(
r#"
UPDATE tts_results
SET
mod_deleted_at = CURRENT_TIMESTAMP,
maybe_mod_user_token = ?
WHERE
token = ?
LIMIT 1
"#,
mod_user_token,
inference_result_token,
)
.execute(mysql_pool)
.await?;
Ok(())
}
pub async fn undelete_tts_inference_result_as_user(
inference_result_token: &str,
mysql_pool: &MySqlPool
) -> AnyhowResult<()> {
let _r = sqlx::query!(
r#"
UPDATE tts_results
SET
user_deleted_at = NULL
WHERE
token = ?
LIMIT 1
"#,
inference_result_token,
)
.execute(mysql_pool)
.await?;
Ok(())
}
pub async fn undelete_tts_inference_result_as_mod(
inference_result_token: &str,
mod_user_token: &str,
mysql_pool: &MySqlPool
) -> AnyhowResult<()> {
let _r = sqlx::query!(
r#"
UPDATE tts_results
SET
mod_deleted_at = NULL,
maybe_mod_user_token = ?
WHERE
token = ?
LIMIT 1
"#,
mod_user_token,
inference_result_token,
)
.execute(mysql_pool)
.await?;
Ok(())
}
@@ -1,5 +1,4 @@
pub mod delete_tts_result_various_scopes;
pub mod edit_tts_result;
pub mod insert_tts_result;
pub mod list_tts_results_query_builder;
pub mod query_tts_result;
pub mod query_tts_result;
@@ -0,0 +1,21 @@
use tokens::tokens::api_keys::ApiKeyToken;
use tokens::tokens::users::UserToken;
/// A user authenticated by either an API key or a web session cookie.
///
/// The two lookup queries (`get_api_or_web_session_user_by_api_key` and
/// `get_api_or_web_session_user_by_session_token`) select congruent user fields so both can
/// produce this one uniform record.
pub struct ApiOrWebSessionUserRecord {
pub user_token: UserToken,
pub username: String,
pub display_name: String,
pub email_address: String,
pub user_role_slug: String,
pub is_banned: bool,
pub can_ban_users: bool,
/// Only set when the lookup was by API key.
pub maybe_api_key_token: Option<ApiKeyToken>,
}
@@ -0,0 +1,70 @@
use std::marker::PhantomData;
use sqlx::{Executor, MySql};
use artcraft_api_keys::ArtcraftApiKey;
use tokens::tokens::api_keys::ApiKeyToken;
use tokens::tokens::users::UserToken;
use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_bool_default_false};
use crate::queries::users::api_or_web_sessions::api_or_web_session_user_record::ApiOrWebSessionUserRecord;
pub struct GetApiOrWebSessionUserByApiKeyArgs<'e, 'c, E>
where
E: 'e + Executor<'c, Database = MySql>,
{
/// The actual API key secret presented by the client.
pub api_key: &'e ArtcraftApiKey,
pub mysql_executor: E,
pub phantom: PhantomData<&'c E>,
}
/// Look up the user behind an API key secret, selecting the same user fields as
/// `get_api_or_web_session_user_by_session_token` so both auth paths share one record shape.
///
/// Returns `Ok(None)` when no row matches, when the matching API key is soft-deleted, or when the
/// owning user is missing or deleted. The full secret is never echoed back.
pub async fn get_api_or_web_session_user_by_api_key<'e, 'c: 'e, E>(
args: GetApiOrWebSessionUserByApiKeyArgs<'e, 'c, E>,
) -> Result<Option<ApiOrWebSessionUserRecord>, sqlx::Error>
where
E: 'e + Executor<'c, Database = MySql>,
{
let result = sqlx::query!(
r#"
SELECT
users.token as `user_token: UserToken`,
users.username,
users.display_name,
users.email_address,
users.user_role_slug,
users.is_banned,
user_roles.can_ban_users,
api_keys.token as `api_key_token: ApiKeyToken`
FROM api_keys
INNER JOIN users
ON users.token = api_keys.owner_user_token
LEFT OUTER JOIN user_roles
ON users.user_role_slug = user_roles.slug
WHERE api_keys.api_key = ?
AND api_keys.maybe_deleted_at IS NULL
AND users.user_deleted_at IS NULL
AND users.mod_deleted_at IS NULL
LIMIT 1
"#,
args.api_key.as_str_be_careful(),
)
.fetch_optional(args.mysql_executor)
.await?;
Ok(result.map(|r| ApiOrWebSessionUserRecord {
user_token: r.user_token,
username: r.username,
display_name: r.display_name,
email_address: r.email_address,
user_role_slug: r.user_role_slug,
is_banned: i8_to_bool(r.is_banned),
can_ban_users: nullable_i8_to_bool_default_false(r.can_ban_users),
maybe_api_key_token: Some(r.api_key_token),
}))
}
@@ -0,0 +1,67 @@
use std::marker::PhantomData;
use sqlx::{Executor, MySql};
use tokens::tokens::users::UserToken;
use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_bool_default_false};
use crate::queries::users::api_or_web_sessions::api_or_web_session_user_record::ApiOrWebSessionUserRecord;
pub struct GetApiOrWebSessionUserBySessionTokenArgs<'e, 'c, E>
where
E: 'e + Executor<'c, Database = MySql>,
{
/// The session token decoded from the session cookie.
pub session_token: &'e str,
pub mysql_executor: E,
pub phantom: PhantomData<&'c E>,
}
/// Look up the user behind a web session token, selecting the same user fields as
/// `get_api_or_web_session_user_by_api_key` so both auth paths share one record shape.
///
/// Returns `Ok(None)` when no row matches, when the session is deleted, or when the owning user
/// is missing or deleted.
pub async fn get_api_or_web_session_user_by_session_token<'e, 'c: 'e, E>(
args: GetApiOrWebSessionUserBySessionTokenArgs<'e, 'c, E>,
) -> Result<Option<ApiOrWebSessionUserRecord>, sqlx::Error>
where
E: 'e + Executor<'c, Database = MySql>,
{
let result = sqlx::query!(
r#"
SELECT
users.token as `user_token: UserToken`,
users.username,
users.display_name,
users.email_address,
users.user_role_slug,
users.is_banned,
user_roles.can_ban_users
FROM user_sessions
INNER JOIN users
ON users.token = user_sessions.user_token
LEFT OUTER JOIN user_roles
ON users.user_role_slug = user_roles.slug
WHERE user_sessions.token = ?
AND user_sessions.deleted_at IS NULL
AND users.user_deleted_at IS NULL
AND users.mod_deleted_at IS NULL
LIMIT 1
"#,
args.session_token,
)
.fetch_optional(args.mysql_executor)
.await?;
Ok(result.map(|r| ApiOrWebSessionUserRecord {
user_token: r.user_token,
username: r.username,
display_name: r.display_name,
email_address: r.email_address,
user_role_slug: r.user_role_slug,
is_banned: i8_to_bool(r.is_banned),
can_ban_users: nullable_i8_to_bool_default_false(r.can_ban_users),
maybe_api_key_token: None,
}))
}
@@ -0,0 +1,3 @@
pub mod api_or_web_session_user_record;
pub mod get_api_or_web_session_user_by_api_key;
pub mod get_api_or_web_session_user_by_session_token;
@@ -1,3 +1,4 @@
pub mod api_or_web_sessions;
pub mod user;
pub mod user_badges;
pub mod user_bookmarks;
@@ -17,10 +17,6 @@ impl RedisCacheKeys {
format!("get_tts_inf2:{}", model_token)
}
pub fn get_unified_queue_stats_endpoint() -> &'static str {
"q_stats_2" // NB: Schema changed, hence _2
}
pub fn session_record_light(session_token: &str) -> String {
format!("ses_light:{}", session_token)
}
@@ -16,7 +16,6 @@ use crate::http_server::common_responses::tag_info::TagInfo;
use crate::http_server::common_responses::user_details_lite::{UserDefaultAvatarInfo, UserDetailsLight};
use crate::http_server::deprecated_endpoints::conversion::enqueue_fbx_to_gltf_handler::*;
use crate::http_server::deprecated_endpoints::engine::create_scene_handler::*;
use crate::http_server::deprecated_endpoints::image_gen::enqueue_image_generation::EnqueueImageGenRequestSuccessResponse;
use crate::http_server::deprecated_endpoints::workflows::enqueue::enqueue_face_fusion_workflow_handler::*;
use crate::http_server::deprecated_endpoints::workflows::enqueue::enqueue_live_portrait_workflow_handler::*;
use crate::http_server::deprecated_endpoints::workflows::enqueue::vst_common::vst_error::*;
@@ -920,7 +919,6 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux
EnqueueFbxToGltfRequestSuccessResponse,
EnqueueGptSovitsModelDownloadRequest,
EnqueueGptSovitsModelDownloadSuccessResponse,
EnqueueImageGenRequestSuccessResponse,
EnqueueLivePortraitCropDimensions,
EnqueueLivePortraitWorkflowRequest,
EnqueueLivePortraitWorkflowSuccessResponse,
@@ -985,15 +983,12 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux
GetProfilePathInfo,
GetPromptPathInfo,
GetPromptSuccessResponse,
GetUnifiedQueueStatsError,
GetUnifiedQueueStatsSuccessResponse,
GetUserRatingResponse,
GetWeightPathInfo,
GetWeightResponse,
GoogleCreateAccountRequest,
GoogleCreateAccountSuccessResponse,
InferTtsError,
InferTtsRequest,
InferTtsSuccessResponse,
InferenceJobStatusResponsePayload,
InferenceJobTokenType,
@@ -1,143 +0,0 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use std::sync::Arc;
use actix_web::web::Path;
use actix_web::{web, HttpRequest, HttpResponse};
use chrono::{DateTime, Utc};
use log::warn;
use mysql_queries::queries::tts::tts_category_assignments::list_assigned_tts_categories_query_builder::ListAssignedTtsCategoriesQueryBuilder;
use crate::state::server_state::ServerState;
use crate::http_server::common_responses::common_web_error::CommonWebError;
// =============== Request ===============
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct ListTtsModelAssignedCategoriesPathInfo {
/// TTS model token
token: String,
}
// =============== Success Response ===============
#[derive(Serialize)]
pub struct ListTtsModelAssignedCategoriesResponse {
pub success: bool,
pub categories: Vec<DisplayCategory>,
}
/// Public-facing and optimized (non-irrelevant fields) category list
/// Used for the main TTS dropdown as well as the TTS edit/tagging UI
#[derive(Serialize)]
pub struct DisplayCategory {
pub category_token: String,
pub model_type: String, // TODO: ENUM
pub maybe_super_category_token: Option<String>,
pub can_directly_have_models: bool,
pub can_have_subcategories: bool,
pub can_only_mods_apply: bool,
pub name: String,
pub maybe_dropdown_name: Option<String>,
// NB: It's okay for non-mods to see this.
pub is_mod_approved: Option<bool>,
pub category_created_at: DateTime<Utc>,
pub category_updated_at: DateTime<Utc>,
pub category_deleted_at: Option<DateTime<Utc>>,
//pub creator_user_token: Option<String>,
//pub creator_username: Option<String>,
//pub creator_display_name: Option<String>,
}
// =============== Error Response ===============
// NB: Not using derive_more::Display since Clion doesn't understand it.
// =============== Handler ===============
pub async fn list_tts_model_assigned_categories_handler(
http_request: HttpRequest,
path: Path<ListTtsModelAssignedCategoriesPathInfo>,
server_state: web::Data<Arc<ServerState>>) -> Result<HttpResponse, CommonWebError>
{
let mut mysql_connection = server_state.mysql_pool.acquire()
.await
.map_err(|e| {
warn!("Could not acquire DB pool: {:?}", e);
CommonWebError::from_error(e)
})?;
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from_error(e)
})?;
// TODO: We don't have any permissions for categories. This is a proxy.
let is_mod = maybe_user_session
.map(|session| session.can_ban_users)
.unwrap_or(false);
let query_builder = ListAssignedTtsCategoriesQueryBuilder::for_model_token(&path.token)
.show_invalid_model_not_allowed_categories(is_mod)
.show_deleted(is_mod)
.show_unapproved(is_mod);
let query_result =
query_builder.perform_query_by_connection(&mut mysql_connection).await;
let results = match query_result {
Ok(results) => results,
Err(err) => {
warn!("DB error: {:?}", err);
return Err(CommonWebError::from_anyhow_error(err));
}
};
let mut categories = results.categories.into_iter()
.map(|c| {
DisplayCategory {
category_token: c.category_token,
model_type: c.model_type,
maybe_super_category_token: c.maybe_super_category_token,
can_directly_have_models: c.can_directly_have_models,
can_have_subcategories: c.can_have_subcategories,
can_only_mods_apply: c.can_only_mods_apply,
name: c.name.clone(),
maybe_dropdown_name:c.maybe_dropdown_name,
is_mod_approved: c.is_mod_approved,
category_created_at: c.category_created_at,
category_updated_at: c.category_updated_at,
category_deleted_at: c.category_deleted_at,
}
})
.collect::<Vec<DisplayCategory>>();
// TODO: Sort by dropdown name too.
categories.sort_by(|c1, c2| c1.name.cmp(&c2.name));
let response = ListTtsModelAssignedCategoriesResponse {
success: true,
categories,
};
let body = serde_json::to_string(&response)
.map_err(CommonWebError::from_error)?;
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(body))
}
@@ -1,4 +1,3 @@
pub mod assign_tts_category;
pub mod list_fully_computed_assigned_tts_categories;
pub mod list_tts_categories;
pub mod list_tts_model_assigned_categories;
@@ -1,75 +0,0 @@
use std::collections::BTreeSet;
#[cfg(test)]
use strum::EnumCount;
#[cfg(test)]
use strum::EnumIter;
use utoipa::ToSchema;
#[cfg_attr(test, derive(EnumIter, EnumCount))]
#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Debug, Deserialize, Serialize, ToSchema)]
pub enum EnqueueImageGenType {
#[serde(rename = "upload_lora")]
UploadLoRA,
#[serde(rename = "upload_sd")]
Checkpoint,
#[serde(rename = "inference")]
Inference,
}
impl EnqueueImageGenType {
pub fn to_str(&self) -> &'static str {
match self {
Self::UploadLoRA => "upload_lora",
Self::Checkpoint => "upload_sd",
Self::Inference => "inference",
}
}
pub fn from_str(value: &str) -> Result<Self, String> {
match value {
"upload_lora" => Ok(Self::UploadLoRA),
"upload_sd" => Ok(Self::Checkpoint),
"inference" => Ok(Self::Inference),
_ => Err(format!("invalid value: {:?}", value)),
}
}
pub fn all_variants() -> BTreeSet<Self> {
BTreeSet::from([
Self::UploadLoRA,
Self::Checkpoint,
Self::Inference
])
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_to_str() {
assert_eq!(EnqueueImageGenType::UploadLoRA.to_str(), "upload_lora");
assert_eq!(EnqueueImageGenType::Checkpoint.to_str(), "upload_sd");
assert_eq!(EnqueueImageGenType::Inference.to_str(), "inference");
}
#[test]
fn test_from_str() {
assert_eq!(EnqueueImageGenType::from_str("upload_lora").unwrap(), EnqueueImageGenType::UploadLoRA);
assert_eq!(EnqueueImageGenType::from_str("upload_sd").unwrap(), EnqueueImageGenType::Checkpoint);
assert_eq!(EnqueueImageGenType::from_str("inference").unwrap(), EnqueueImageGenType::Inference);
assert!(EnqueueImageGenType::from_str("invalid").is_err());
}
#[test]
fn test_all_variants() {
let variants = EnqueueImageGenType::all_variants();
assert_eq!(variants.len(), 3);
assert!(variants.contains(&EnqueueImageGenType::UploadLoRA));
assert!(variants.contains(&EnqueueImageGenType::Checkpoint));
assert!(variants.contains(&EnqueueImageGenType::Inference));
}
}
@@ -1,560 +0,0 @@
#![forbid(unused_mut)]
use std::fmt::Debug;
use std::fmt::{Display, Formatter};
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::web::Json;
use actix_web::{web, HttpRequest, HttpResponse};
use log::error;
use log::warn;
use serde::Deserialize;
use serde::Serialize;
use utoipa::ToSchema;
use web::Data;
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
use http_server_common::request::get_request_header_optional::get_request_header_optional;
use http_server_common::request::get_request_ip::get_request_ip;
use mysql_queries::payloads::generic_inference_args::generic_inference_args::{
GenericInferenceArgs,
InferenceCategoryAbbreviated,
PolymorphicInferenceArgs,
};
use mysql_queries::payloads::generic_inference_args::inner_payloads::image_generation_payload::StableDiffusionArgs;
use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{
insert_generic_inference_job,
InsertGenericInferenceArgs,
};
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::model_weights::ModelWeightToken;
use tokens::tokens::users::UserToken;
use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::state::server_state::ServerState;
/// This is the number of images (batch size) to generate for each request.
/// We should allow all users to have multiple images generated at once as this
/// is what other providers do.
const MINIMUM_IMAGE_COUNT : u32 = 4;
/// The maximum number of images (batch size) to generate for each request.
const MAXIMUM_IMAGE_COUNT : u32 = 8;
/// Debug requests can get routed to special "debug-only" workers, which can
/// be used to trial new code, run debugging, etc.
const DEBUG_HEADER_NAME: &str = "enable-debug-mode";
/// The routing tag header can send workloads to particular k8s hosts.
/// This is useful for catching the live logs or intercepting the job.
const ROUTING_TAG_HEADER_NAME: &str = "routing-tag";
#[derive(Debug, Deserialize, Clone, Copy, Eq, PartialEq)]
pub enum TypeOfInference {
#[serde(rename = "lora")]
Lora,
#[serde(rename = "model")]
Model,
#[serde(rename = "inference")]
Inference
}
impl Display for TypeOfInference {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
match self {
TypeOfInference::Lora => write!(f, "lora"),
TypeOfInference::Model => write!(f, "model"),
TypeOfInference::Inference => write!(f, "inference"),
}
}
}
pub fn is_valid_string(input: &str) -> bool {
let valid_strings = [
"DPM++ 2M Karras",
"DPM++ SDE Karras",
"DPM++ 2M SDE Exponential",
"DPM++ 2M SDE Karras",
"Euler a",
"Euler",
"LMS",
"Heun",
"DPM2",
"DPM2 a",
"DPM++ 2S a",
"DPM++ 2M",
"DPM++ SDE",
"DPM++ 2M SDE",
"DPM++ 2M SDE Heun",
"DPM++ 2M SDE Heun Karras",
"DPM++ 2M SDE Heun Exponential",
"DPM++ 3M SDE",
"DPM++ 3M SDE Karras",
"DPM++ 3M SDE Exponential",
"DPM fast",
"DPM adaptive",
"LMS Karras",
"DPM2 Karras",
"DPM2 a Karras",
"DPM++ 2S a Karras",
];
valid_strings.contains(&input)
}
#[derive(Deserialize, Default, ToSchema)]
pub struct EnqueueImageGenRequest {
pub uuid_idempotency_token: String,
pub maybe_image_source: Option<String>,
pub maybe_sd_model_token: Option<String>,
pub maybe_lora_model_token: Option<String>,
pub maybe_prompt: Option<String>,
pub maybe_n_prompt: Option<String>,
pub maybe_seed: Option<i64>,
pub maybe_width: Option<u32>,
pub maybe_height: Option<u32>,
pub maybe_sampler: Option<String>,
pub maybe_upload_path: Option<String>,
pub maybe_lora_upload_path: Option<String>,
pub maybe_cfg_scale: Option<u32>,
pub maybe_number_of_samples: Option<u32>,
pub maybe_batch_count: Option<u32>,
pub maybe_name: Option<String>,
pub maybe_description: Option<String>,
pub maybe_version: Option<u32>,
// Optional cover image on LoRA or SD upload
pub maybe_cover_image_media_file_token: Option<MediaFileToken>,
}
#[derive(Serialize, ToSchema)]
pub struct EnqueueImageGenRequestSuccessResponse {
pub success: bool,
pub inference_job_token: InferenceJobToken,
}
// Implementation for enqueuing a TTS request
// Reference enqueue_infer_tts_handler.rs for checks: rate limiting / user sessions
// insert generic inference job.rs
// Need to convert it to generic inference job.
#[utoipa::path(
post,
path = "/inference/enqueue_image_gen/",
responses(
(
status = 200,
description = "Enqueue TTS generically",
body = EnqueueImageGenRequestSuccessResponse,
),
(status = 400, description = "Bad input", body = CommonWebError),
(status = 401, description = "Not authorized", body = CommonWebError),
(status = 429, description = "Rate limited", body = CommonWebError),
(status = 500, description = "Server error", body = CommonWebError)
),
params(("request" = EnqueueImageGenRequest, description = "Payload for TTS Request"))
)]
pub async fn enqueue_image_generation_request(
http_request: HttpRequest,
request: Json<EnqueueImageGenRequest>,
server_state: Data<Arc<ServerState>>
) -> Result<HttpResponse, CommonWebError> {
let inference_mode = inference_mode_from_http_request(&http_request)
.ok_or(CommonWebError::BadInputWithSimpleMessage("Invalid request".to_string()))?;
validate_request(&request, inference_mode)?;
// TODO: Brandon need to figure out premium vs not premium
let mut maybe_user_token: Option<UserToken> = None;
let visbility = enums::common::visibility::Visibility::Public;
let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| {
warn!("MySql pool error: {:?}", err);
CommonWebError::from_error(err)
})?;
// ==================== USER SESSION ==================== //
let maybe_user_session = server_state.session_checker
.maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection).await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from_error(e)
})?;
if let Some(user_session) = maybe_user_session.as_ref() {
maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token));
}
// ==================== PLANS ==================== //
// Plan should handle "first anonymous use" and "investor" cases.
let plan = get_correct_plan_for_session(
server_state.server_environment,
maybe_user_session.as_ref()
);
// Separate priority for animation.
let priority_level = plan.web_vc_base_priority_level();
// ==================== DEBUG MODE + ROUTING TAG ==================== //
let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME).is_some();
let maybe_routing_tag = get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME).map(
|routing_tag| routing_tag.trim().to_string()
);
// ==================== BANNED USERS ==================== //
if let Some(ref user) = maybe_user_session {
if user.role.is_banned {
return Err(CommonWebError::NotAuthorized);
}
}
// DETECT premium user and queue
// ==================== RATE LIMIT ==================== //
let rate_limiter = match maybe_user_session {
None => &server_state.redis_rate_limiters.logged_out,
Some(ref _user) => &server_state.redis_rate_limiters.logged_in,
};
if let Err(_err) = rate_limiter.rate_limit_request(&http_request).await {
return Err(CommonWebError::TooManyRequests);
}
// Get up IP address
let ip_address = get_request_ip(&http_request);
// Check the inference args to make sure everything is all there for upload loRA / model or standard inference
// ==================== HANDLE IDEMPOTENCY ==================== //
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection)
.await
.map_err(|err| {
error!("Error inserting idempotency token: {:?}", err);
CommonWebError::BadInputWithSimpleMessage("invalid idempotency token".to_string())
})?;
// ==================== INFERENCE ARGS ==================== //
let width = request.maybe_width.unwrap_or(512);
let height = request.maybe_height.unwrap_or(512);
let width = match width {
0..=512 => 512,
513..=768 => 768,
769..=1024 => 1024,
_ => 1024,
};
let height = match height {
0..=512 => 512,
513..=768 => 768,
769..=1024 => 1024,
_ => 1024,
};
let cfg_scale = match request.maybe_cfg_scale {
Some(val) => if val > 32 { 32 } else { val }
None => 7,
};
let number_of_samples = match request.maybe_number_of_samples {
Some(val) => if val > 128 { 128 } else { val }
None => 20,
};
let batch_count = match request.maybe_batch_count {
None => MINIMUM_IMAGE_COUNT, // NB: Default to "3" images for everyone
Some(0...MINIMUM_IMAGE_COUNT) => MINIMUM_IMAGE_COUNT,
Some(val) => if val > MAXIMUM_IMAGE_COUNT { MAXIMUM_IMAGE_COUNT } else { val }
};
let sampler = match request.maybe_sampler.clone() {
Some(val) => if is_valid_string(&val) { val.clone() } else { String::from("Euler a") }
None => String::from("Euler a"),
};
let sd_weight_token = ModelWeightToken(
request.maybe_sd_model_token.clone().unwrap_or_default()
);
let lora_token = ModelWeightToken(request.maybe_lora_model_token.clone().unwrap_or_default());
let n_prompt = request.maybe_n_prompt.clone().unwrap_or_default();
// we can only do 1 upload type at a time.
// if we are uploading a model.
let mut maybe_sd_upload_url = None;
let mut maybe_lora_upload_url = None;
let both_fields = (
request.maybe_upload_path.as_deref(),
request.maybe_lora_upload_path.as_deref()
);
match both_fields {
(Some(_), Some(_)) => {
return Err(CommonWebError::BadInputWithSimpleMessage("Can't upload both lora and model".to_string()));
}
(Some(sd_url), None) => {
maybe_sd_upload_url = Some(sd_url.to_string());
}
(None, Some(lora_url)) => {
maybe_lora_upload_url = Some(lora_url.to_string());
}
_ => {}, // No-op for inference
}
// NB: This is done to populate the new top-level field in the jobs table.
// We won't read this yet, but in the meantime it can serve as analytics, and in the future
// we can switch to it.
let maybe_either_download_url = maybe_sd_upload_url.as_deref()
.or(maybe_lora_upload_url.as_deref())
.map(|s| s.to_string());
let seed = request.maybe_seed.unwrap_or(-1);
let mut version : u32 = 0;
let type_of_inference = match inference_mode {
TypeOfInference::Lora => "lora",
TypeOfInference::Model => "model",
TypeOfInference::Inference => "inference",
};
if let Some(s) = request.maybe_version {
version = s;
}
let inference_args = StableDiffusionArgs {
maybe_sd_model_token: Some(sd_weight_token),
maybe_lora_model_token: Some(lora_token),
maybe_prompt: request.maybe_prompt.clone(),
maybe_n_prompt: request.maybe_n_prompt.clone(),
maybe_seed: Some(seed),
maybe_upload_path: maybe_sd_upload_url,
maybe_lora_upload_path: maybe_lora_upload_url,
type_of_inference: type_of_inference.to_string(),
maybe_cfg_scale: Some(cfg_scale),
maybe_number_of_samples: Some(number_of_samples),
maybe_batch_count: Some(batch_count),
maybe_width: Some(width),
maybe_height: Some(height),
maybe_sampler: Some(sampler),
maybe_description: request.maybe_description.clone(),
maybe_name: request.maybe_name.clone(),
maybe_version: Some(version),
};
// create the inference args here
let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request);
// create the job record here!
let query_result = insert_generic_inference_job(InsertGenericInferenceArgs {
uuid_idempotency_token: &request.uuid_idempotency_token,
job_type: InferenceJobType::StableDiffusion,
maybe_product_category: None, // This is not a product anymore
inference_category: InferenceCategory::ImageGeneration,
maybe_model_type: Some(InferenceModelType::StableDiffusion), // NB: Model is static during inference
maybe_model_token: None, // NB: Model is static during inference
maybe_input_source_token: None,
maybe_input_source_token_type: None,
maybe_download_url: maybe_either_download_url.as_deref(),
maybe_cover_image_media_file_token: request.maybe_cover_image_media_file_token.as_ref(),
maybe_raw_inference_text: None,
maybe_max_duration_seconds: None,
maybe_inference_args: Some(GenericInferenceArgs {
inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration),
args: Some(PolymorphicInferenceArgs::Ig(inference_args)),
}),
maybe_creator_user_token: maybe_user_token.as_ref(),
maybe_avt_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility: visbility,
priority_level,
requires_keepalive: false, //reverse ... TODO fix this. we set it base on account is premium or not ...
is_debug_request,
maybe_routing_tag: maybe_routing_tag.as_deref(),
mysql_pool: &server_state.mysql_pool,
}).await;
let job_token = match query_result {
Ok((job_token, _id)) => job_token,
Err(err) => {
warn!("New generic inference job creation DB error: {:?}", err);
if err.had_duplicate_idempotency_token() {
return Err(CommonWebError::BadInputWithSimpleMessage("Duplicate idempotency token".to_string()));
}
return Err(CommonWebError::from_error(err));
}
};
let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse {
success: true,
inference_job_token: job_token,
};
let body = serde_json
::to_string(&response)
?;
// Error handling 101 rust result type returned like so.
Ok(HttpResponse::Ok().content_type("application/json").body(body))
}
fn validate_request(
request: &Json<EnqueueImageGenRequest>,
inference_mode: TypeOfInference
) -> Result<(), CommonWebError> {
let mut requires_name = false;
let mut requires_description = false;
let mut requires_sd_model_token = false;
let mut requires_lora_upload_path = false;
let mut requires_upload_path = false;
match inference_mode {
TypeOfInference::Lora => {
requires_name = true;
requires_description = true;
requires_lora_upload_path = true;
}
TypeOfInference::Model => {
requires_name = true;
requires_description = true;
requires_upload_path = true;
}
TypeOfInference::Inference => {
requires_sd_model_token = true;
}
}
if requires_name && request.maybe_name.is_none() {
return Err(CommonWebError::BadInputWithSimpleMessage("Missing Model / Lora Name".to_string()));
}
if requires_description && request.maybe_description.is_none() {
return Err(CommonWebError::BadInputWithSimpleMessage("Missing Model / Lora Description".to_string()));
}
if requires_sd_model_token && request.maybe_sd_model_token.is_none() {
return Err(CommonWebError::BadInputWithSimpleMessage("Missing Model Token".to_string()));
}
if requires_lora_upload_path && request.maybe_lora_upload_path.is_none() {
return Err(CommonWebError::BadInputWithSimpleMessage("Missing Lora Upload Path".to_string()));
}
if requires_upload_path && request.maybe_upload_path.is_none() {
return Err(CommonWebError::BadInputWithSimpleMessage("Missing Model Upload Path".to_string()));
}
Ok(())
}
fn inference_mode_from_http_request(http_request: &HttpRequest) -> Option<TypeOfInference> {
inference_mode_from_url_path(http_request.path())
}
fn inference_mode_from_url_path(url_path: &str) -> Option<TypeOfInference> {
let last_segment = url_path.split("/").last();
match last_segment {
Some("lora") => Some(TypeOfInference::Lora),
Some("model") => Some(TypeOfInference::Model),
Some("inference") => Some(TypeOfInference::Inference),
_ => None
}
}
#[cfg(test)]
mod tests {
use crate::http_server::deprecated_endpoints::image_gen::enqueue_image_generation::{inference_mode_from_url_path, TypeOfInference};
#[test]
fn test_url_paths() {
// Valid routes
assert_eq!(inference_mode_from_url_path("/v1/image_gen/enqueue/inference"), Some(TypeOfInference::Inference));
assert_eq!(inference_mode_from_url_path("/v1/image_gen/upload/model"), Some(TypeOfInference::Model));
assert_eq!(inference_mode_from_url_path("/v1/image_gen/upload/lora"), Some(TypeOfInference::Lora));
// Non-routes
assert_eq!(inference_mode_from_url_path(""), None);
assert_eq!(inference_mode_from_url_path("/v1/image_gen/enqueue/foo"), None);
assert_eq!(inference_mode_from_url_path("/v1/image_gen/upload/foo"), None);
}
}
// with LoRA
// http://127.0.0.1:12345/v1/image_gen/inference/enqueue_image_gen
// {
// "uuid_idempotency_token": "12",
// "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8",
// "maybe_lora_model_token": "weight_t7gz78fjg27m0wtw6r33gafxs",
// "maybe_prompt": "raiden mei a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed, highest quality, semi nude, masterpiece, painted",
// "maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.",
// "maybe_n_prompt": "nsfw, black and white, low quality, pixelated",
// "maybe_seed": -1,
// "maybe_width": 1024,
// "maybe_height": 1024,
// "maybe_sampler": "DPM++ 2M SDE Heun",
// "maybe_cfg_scale": 7,
// "maybe_number_of_samples": 64,
// "maybe_batch_count": 4,
// }
// without LoRA
// {
// "uuid_idempotency_token": "132141",
// "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8",
// "maybe_prompt": "raiden mei a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed, highest quality, semi nude, masterpiece, painted",
// "maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.",
// "maybe_n_prompt": "nsfw, black and white, low quality, pixelated",
// "maybe_seed": -1,
// "maybe_width": 512,
// "maybe_height": 512,
// "maybe_sampler": "DPM++ 2M SDE Heun",
// "maybe_cfg_scale": 8,
// "maybe_number_of_samples": 64,
// "maybe_batch_count": 4
// }
// http://127.0.0.1:12345/v1/image_gen/upload/lora
// {
// "uuid_idempotency_token": "12",
// "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing",
// "maybe_name":"some_name",
// "maybe_description":"some_description"
// }
// http://127.0.0.1:12345/v1/image_gen/upload/model
// {
// "uuid_idempotency_token": "13",
// "maybe_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing",
// "maybe_name":"some_name",
// "maybe_description":"some_description"
// }
// {
// "uuid_idempotency_token": "123",
// "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing",
// "maybe_name":"some_name",
// "maybe_description":"some_description"
// }
@@ -1,18 +0,0 @@
{
"uuid_idempotency_token": "123",
"type_of_inference": "checkpoint",
"maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8",
"maybe_lora_model_token": "weight_t7gz78fjg27m0wtw6r33gafxs",
"maybe_prompt": "a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed.",
"maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.",
"maybe_n_prompt": "nsfw, black and white, low quality, pixelated",
"maybe_seed": "-1",
"maybe_width": "512",
"maybe_height": "512",
"maybe_sampler": "DPM++ 2M SDE Heun",
"maybe_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing",
"maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing",
"maybe_cfg_scale": "7",
"maybe_number_of_samples": "32",
"maybe_batch_count": "4"
}
@@ -1,2 +0,0 @@
pub mod enqueue_image_generation;
pub mod enqueue_generation_type;
@@ -1,199 +0,0 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::{web, HttpRequest, HttpResponse};
use log::{debug, error, warn};
use sqlx::MySqlPool;
use crate::http_server::common_responses::user_avatars::default_avatar_color_from_username::default_avatar_color_from_username;
use crate::http_server::common_responses::user_avatars::default_avatar_from_username::default_avatar_from_username;
use crate::http_server::common_responses::user_details_lite::{UserDefaultAvatarInfo, UserDetailsLight};
use errors::AnyhowResult;
use mysql_queries::queries::tts::stats::calculate_tts_model_leaderboard::calculate_tts_model_leaderboard;
use tokens::tokens::users::UserToken;
use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error;
use crate::state::server_state::ServerState;
#[derive(Serialize)]
pub struct LeaderboardResponse {
success: bool,
tts_leaderboard: Vec<LeaderboardRow>,
w2l_leaderboard: Vec<LeaderboardRow>,
}
#[derive(Clone, Serialize)]
pub struct LeaderboardRow {
pub user: UserDetailsLight,
pub uploaded_count: i64,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub creator_user_token: String,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub username: String,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub display_name: String,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub gravatar_hash: String,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub default_avatar_index: u8,
#[deprecated(note="switch to the user field (type UserDetailsLight)")]
pub default_avatar_color_index: u8,
}
#[derive(Serialize, Debug)]
pub struct LeaderboardErrorResponse {
pub success: bool,
pub error_type: LeaderboardErrorType,
pub error_message: String,
}
#[derive(Copy, Clone, Debug, Serialize)]
pub enum LeaderboardErrorType {
ServerError,
}
impl LeaderboardErrorResponse {
fn server_error() -> Self {
Self {
success: false,
error_type: LeaderboardErrorType::ServerError,
error_message: "server error".to_string()
}
}
}
impl ResponseError for LeaderboardErrorResponse {
fn status_code(&self) -> StatusCode {
match self.error_type {
LeaderboardErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn error_response(&self) -> HttpResponse {
serialize_as_json_error(self)
}
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
impl fmt::Display for LeaderboardErrorResponse {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self.error_type)
}
}
pub async fn leaderboard_handler(
_http_request: HttpRequest,
server_state: web::Data<Arc<ServerState>>
) -> Result<HttpResponse, LeaderboardErrorResponse> {
let maybe_cached = server_state.caches.ephemeral.leaderboard.grab_copy_without_bump_if_unexpired()
.map_err(|e| {
error!("error consulting cache: {:?}", e);
LeaderboardErrorResponse::server_error()
})?;
let leaderboard_info = match maybe_cached {
Some(leaderboard) => {
leaderboard
}
None => {
debug!("populating leaderboard from database");
let leaderboard_query_result = query_leaderboard(
&server_state.mysql_pool
).await;
match leaderboard_query_result {
// If the database misbehaves (eg. DDoS), let's stop spamming it.
// We'll attempt to read the old value from the cache and keep going.
Err(err) => {
warn!("error querying database / inserting into cache: {:?}", err);
let maybe_cached = server_state.caches.ephemeral.leaderboard.grab_even_expired_and_bump()
.map_err(|err| {
error!("error consulting cache (even expired): {:?}", err);
LeaderboardErrorResponse::server_error()
})?;
maybe_cached.ok_or_else(|| {
error!("error querying database and subsequently reading cache: {:?}", err);
LeaderboardErrorResponse::server_error()
})?
}
// Happy path...
Ok(leaderboard_info) => {
server_state.caches.ephemeral.leaderboard.store_copy(&leaderboard_info)
.map_err(|e| {
error!("error storing cache: {:?}", e);
LeaderboardErrorResponse::server_error()
})?;
leaderboard_info
}
}
}
};
let response = LeaderboardResponse {
success: true,
tts_leaderboard: leaderboard_info.tts_leaderboard,
w2l_leaderboard: vec![],
};
let body = serde_json::to_string(&response)
.map_err(|_e| LeaderboardErrorResponse::server_error())?;
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(body))
}
#[derive(Clone)]
pub struct LeaderboardInfo {
tts_leaderboard: Vec<LeaderboardRow>,
}
async fn query_leaderboard(mysql_pool: &MySqlPool) -> AnyhowResult<LeaderboardInfo> {
let mut mysql_connection = mysql_pool.acquire().await?;
let tts_results = calculate_tts_model_leaderboard(&mut mysql_connection)
.await?
.into_iter()
.map(|record| LeaderboardRow {
creator_user_token: record.creator_user_token.clone(),
username: record.username.to_string(),
display_name: record.display_name.to_string(),
gravatar_hash: record.gravatar_hash.to_string(),
default_avatar_index: default_avatar_from_username(&record.username),
default_avatar_color_index: default_avatar_color_from_username(&record.username),
user: UserDetailsLight {
user_token: UserToken::new_from_str(&record.creator_user_token),
username: record.username.to_string(),
display_name: record.display_name,
gravatar_hash: record.gravatar_hash,
default_avatar: UserDefaultAvatarInfo::from_username(&record.username),
},
uploaded_count: record.uploaded_count,
})
.collect();
Ok(LeaderboardInfo {
tts_leaderboard: tts_results,
})
}
@@ -1 +0,0 @@
pub mod get_leaderboard;
@@ -5,8 +5,6 @@ pub mod conversion;
pub mod engine;
pub mod events;
pub mod flags;
pub mod image_gen;
pub mod leaderboard;
pub mod media_uploads;
pub mod moderation;
pub mod stubs;
@@ -24,6 +24,7 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use url_utils::extension::extract_extension_from_url::{extract_extension_from_url, ExtractExtensions};
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::user_lookup::api_or_web_session::require_api_or_web_session::require_api_or_web_session;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use crate::util::http_download_url_to_bytes::http_download_url_to_bytes;
@@ -58,19 +59,12 @@ pub async fn create_character_handler(
.acquire()
.await?;
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from(e)
})?;
let user_session = match maybe_user_session {
Some(session) => session,
None => return Err(CommonWebError::NotAuthorized),
};
let user_session = require_api_or_web_session(
&http_request,
&server_state.session_checker,
&server_state.avt_cookie_manager,
&mut *mysql_connection,
).await?;
let user_token = &user_session.user_token;
let ip_address = get_request_ip(&http_request);
@@ -194,7 +188,7 @@ pub async fn create_character_handler(
kinovi_character_id: &gen_result.character_id,
maybe_model_type: None,
maybe_creator_user_token: Some(user_token),
maybe_avt_token: None,
maybe_avt_token: user_session.maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility: Visibility::Hidden,
mysql_executor: &mut *transaction,
@@ -10,7 +10,7 @@ use mysql_queries::queries::characters::delete_character::delete_character;
use mysql_queries::queries::characters::get_character_by_token_including_deleted::get_character_by_token_including_deleted;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::user_lookup::user_session::require_user_session::require_user_session;
use crate::http_server::user_lookup::api_or_web_session::require_api_or_web_session::require_api_or_web_session;
use crate::state::server_state::ServerState;
/// Delete a character.
@@ -40,7 +40,12 @@ pub async fn delete_character_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
let user_session = require_user_session(&http_request, &server_state.session_checker, &mut *mysql_connection).await?;
let user_session = require_api_or_web_session(
&http_request,
&server_state.session_checker,
&server_state.avt_cookie_manager,
&mut *mysql_connection,
).await?;
let user_token = &user_session.user_token;
let is_mod = user_session.is_mod();
@@ -11,7 +11,7 @@ use seedance2pro_client::creds::seedance2pro_session::Seedance2ProSession;
use seedance2pro_client::requests::update_character::update_character::{update_character, UpdateCharacterArgs};
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::user_lookup::user_session::require_user_session::require_user_session;
use crate::http_server::user_lookup::api_or_web_session::require_api_or_web_session::require_api_or_web_session;
use crate::state::server_state::ServerState;
/// Edit a character's name or description.
@@ -38,7 +38,12 @@ pub async fn edit_character_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
let user_session = require_user_session(&http_request, &server_state.session_checker, &mut *mysql_connection).await?;
let user_session = require_api_or_web_session(
&http_request,
&server_state.session_checker,
&server_state.avt_cookie_manager,
&mut *mysql_connection,
).await?;
let user_token = &user_session.user_token;
let is_mod = user_session.is_mod();
@@ -16,7 +16,7 @@ use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::common_responses::media::media_domain::MediaDomain;
use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
use crate::http_server::user_lookup::user_session::require_user_session::require_user_session;
use crate::http_server::user_lookup::api_or_web_session::require_api_or_web_session::require_api_or_web_session;
use crate::state::server_state::ServerState;
/// List characters for the current session.
@@ -39,7 +39,12 @@ pub async fn list_characters_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
let user_session = require_user_session(&http_request, &server_state.session_checker, &mut *mysql_connection).await?;
let user_session = require_api_or_web_session(
&http_request,
&server_state.session_checker,
&server_state.avt_cookie_manager,
&mut *mysql_connection,
).await?;
let user_token = &user_session.user_token;
@@ -1,12 +1,22 @@
use log::warn;
use sqlx::{Executor, MySql};
use artcraft_router::api::router_provider::RouterProvider;
use enums::by_table::debug_logs::debug_log_level::DebugLogLevel;
use enums::by_table::debug_logs::debug_log_type::DebugLogType;
use mysql_queries::queries::debug_logs::insert_debug_log::{insert_debug_log, InsertDebugLogArgs};
use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken;
use tokens::tokens::users::UserToken;
/// Request-scoped identity for debug log inserts, so deep pipeline code can
/// log without threading four separate arguments.
pub struct GenerationDebugLogContext<'a> {
pub event_token: &'a DebugLogEventToken,
pub user_token: &'a UserToken,
pub ip_address: &'a str,
pub request_url: &'a str,
}
/// Best-effort info-level debug log of an inbound generation HTTP request.
/// Never fails the request — insert errors are logged and swallowed.
pub async fn insert_generation_request_debug_log<'e, 'c: 'e, E>(
@@ -60,3 +70,47 @@ pub async fn insert_generation_failure_debug_log<'e, 'c: 'e, E>(
warn!("Failed to insert generation failure debug log: {:?}", err);
}
}
/// The debug log type for an outbound request to the given provider, or
/// `None` for providers we don't capture.
pub fn provider_request_debug_log_type(provider: RouterProvider) -> Option<DebugLogType> {
match provider {
RouterProvider::Fal => Some(DebugLogType::FalRequest),
RouterProvider::GrokApi => Some(DebugLogType::GrokApiRequest),
RouterProvider::Seedance2Pro => Some(DebugLogType::KinoviRequest),
// First-party: the HttpRequest log already captures the user request.
RouterProvider::Artcraft => None,
// No DebugLogType variant yet; add one if GmiCloud routing is revived.
RouterProvider::GmiCloud => None,
}
}
/// Best-effort info-level debug log of an outbound provider generation
/// request. MUST be inserted during the pre-request DB phase — on the same
/// open connection as the session/billing/HTTP-request writes, BEFORE that
/// connection is released and the external provider call is made — so the
/// outbound payload is captured even when the upload/enqueue fails, without
/// ever touching the pool.
/// Never fails the request — insert errors are logged and swallowed.
pub async fn insert_provider_request_debug_log<'e, 'c: 'e, E>(
context: &'e GenerationDebugLogContext<'e>,
debug_log_type: DebugLogType,
outbound_request_debug: &'e str,
mysql_executor: E,
) where
E: 'c + Executor<'c, Database = MySql>,
{
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(context.event_token),
maybe_creator_user_token: Some(context.user_token),
debug_log_type,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(context.ip_address),
maybe_url: Some(context.request_url),
message: outbound_request_debug,
mysql_executor,
phantom: Default::default(),
}).await {
warn!("Failed to insert {} debug log: {:?}", debug_log_type, err);
}
}
@@ -8,6 +8,7 @@ use sqlx::Acquire;
use artcraft_api_defs::omni_api::generate_requests::omni_api_image_generate_request::OmniApiImageGenerateRequest;
use artcraft_api_defs::omni_gen::generate_response::omni_gen_image_generate_response::OmniGenImageGenerateResponse;
use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse;
use enums::by_table::debug_logs::debug_log_level::DebugLogLevel;
use enums::by_table::debug_logs::debug_log_type::DebugLogType;
use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType;
use enums::by_table::prompts::prompt_type::PromptType;
@@ -16,7 +17,6 @@ use enums::common::generation::common_model_type::CommonModelType;
use enums::common::generation_provider::GenerationProvider;
use enums::common::platform_type::PlatformType;
use http_server_common::request::get_request_ip::get_request_ip;
use enums::by_table::debug_logs::debug_log_level::DebugLogLevel;
use mysql_queries::queries::debug_logs::insert_debug_log::{insert_debug_log, InsertDebugLogArgs};
use mysql_queries::queries::generic_inference::api_providers::seedance2pro::insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token::KinoviVersion;
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
@@ -28,6 +28,7 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::GenerationDebugLogContext;
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
use crate::http_server::endpoints::omni_api::generate::image::check_request::check_request;
use crate::http_server::endpoints::omni_api::generate::image::hydrate_to_router_request::hydrate_to_router_request;
@@ -36,8 +37,8 @@ use crate::http_server::endpoints::omni_api::generate::image::insert_db_job::ins
use crate::http_server::endpoints::omni_api::generate::image::insert_db_job::insert_seedance2pro_jobs::{insert_seedance2pro_jobs, InsertSeedance2proJobsArgs};
use crate::http_server::endpoints::omni_api::generate::image::insert_db_job::shared_job_args::SharedJobArgs;
use crate::http_server::endpoints::omni_api::generate::image::pipeline_v2::run_pipeline_v2::{run_pipeline_v2, RunPipelineV2Args};
use crate::http_server::user_lookup::user_session::session_utils::lookup::user_session_feature_flags::UserSessionFeatureFlags;
use crate::http_server::user_lookup::api_keys::require_api_key_user::require_api_key_user;
use crate::http_server::user_lookup::user_session::session_utils::lookup::user_session_feature_flags::UserSessionFeatureFlags;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use crate::util::lookup::lookup_media_files_as_cdn_url_list_and_map::lookup_media_files_as_cdn_url_list_and_map;
@@ -152,16 +153,25 @@ pub async fn omni_api_image_generate_handler(
// ==================== PIPELINE ==================== //
// NB: Release the pooled DB connection before the (slow, external) generation call so we don't
// hold a pool slot idle while waiting on the provider — that's what starves the pool and causes
// PoolTimedOut on unrelated endpoints. We re-acquire below to write the result.
drop(mysql_connection);
// NB: The pipeline takes over the connection for its remaining pre-request DB writes (billing,
// outbound provider request debug log) and releases it before the (slow, external) generation
// call — holding a pool slot across that call is what starves the pool and causes PoolTimedOut
// on unrelated endpoints. We re-acquire below to write the result.
let debug_log_context = GenerationDebugLogContext {
event_token: &debug_log_event_token,
user_token,
ip_address: &ip_address,
request_url: &request_url,
};
let pipeline_result = run_pipeline_v2(RunPipelineV2Args {
router_builder: &router_builder,
server_state: &server_state,
user_token,
resolved_media: &resolved_media,
debug_log_context: &debug_log_context,
mysql_connection,
}).await;
// ==================== DEBUG LOG: PIPELINE ERROR ==================== //
@@ -191,25 +201,8 @@ pub async fn omni_api_image_generate_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
// ==================== DEBUG LOG: FAL REQUEST ==================== //
if let GenerateImageResponse::Fal(ref fal_payload) = pipeline_result.response {
if let Some(ref outbound_request) = fal_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::FalRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Fal request debug log: {:?}", err);
}
}
}
// NB: Outbound provider requests (Fal/Grok/Kinovi) are debug-logged inside
// the pipeline BEFORE the send, so the payload is captured even on failure.
// ==================== WRITE RESULT ==================== //
@@ -17,6 +17,9 @@ use tokens::tokens::users::UserToken;
use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::{
insert_provider_request_debug_log, provider_request_debug_log_type, GenerationDebugLogContext,
};
use crate::http_server::endpoints::omni_api::generate::image::pipeline_result::ImagePipelineResult;
use crate::http_server::endpoints::omni_api::shared_utils::map_seedance2pro_router_error::map_router_error_to_web_error;
use crate::state::server_state::ServerState;
@@ -27,6 +30,11 @@ pub struct RunPipelineV2Args<'a> {
pub server_state: &'a ServerState,
pub user_token: &'a UserToken,
pub resolved_media: &'a MediaFilesAsCdnUrlListAndMap,
pub debug_log_context: &'a GenerationDebugLogContext<'a>,
/// The handler's open connection. The pipeline uses it for its remaining
/// pre-request DB writes (billing, outbound-request debug log) and releases
/// it BEFORE the external provider call.
pub mysql_connection: sqlx::pool::PoolConnection<sqlx::MySql>,
}
// NB: This pipeline does an external generation call (`finalize_and_generate`) that can take many
@@ -41,6 +49,8 @@ pub async fn run_pipeline_v2(
server_state,
user_token,
resolved_media,
debug_log_context,
mut mysql_connection,
} = args;
let hydrated_builder = apply_hydrated_media_inputs(
@@ -57,24 +67,36 @@ pub async fn run_pipeline_v2(
let apriori_job_token = InferenceJobToken::generate();
let maybe_wallet_ledger_entry_token = if cost > 0 {
// NB: Short-lived connection — released before the external call below.
let mut billing_connection = server_state.mysql_pool.acquire().await
.map_err(|err| {
warn!("Failed to acquire MySQL connection for image billing: {:?}", err);
CommonWebError::from_error(err)
})?;
// Billed on the handler's connection (same pre-request DB phase).
let deduction = attempt_wallet_deduction_else_common_web_error(
user_token,
Some(apriori_job_token.as_str()),
cost,
&mut billing_connection,
&mut mysql_connection,
).await?;
Some(deduction.ledger_entry_token)
} else {
None
};
// NB: No pooled DB connection is held across this external generation call.
// Debug-log the outbound provider request BEFORE the send — still on the
// handler's connection — so the payload is captured even when the
// upload/enqueue fails.
if let Some(debug_log_type) = provider_request_debug_log_type(draft_or_request.get_provider()) {
insert_provider_request_debug_log(
debug_log_context,
debug_log_type,
&format!("{:#?}", draft_or_request),
&mut *mysql_connection,
).await;
}
// NB: Done with pre-request DB writes. Release the pooled connection before
// the (slow, external) provider call — holding it across that call is what
// starves the pool and causes PoolTimedOut. No pooled DB connection is held
// across the external generation call below.
drop(mysql_connection);
let response = finalize_and_generate(draft_or_request, server_state, resolved_media).await?;
Ok(ImagePipelineResult {
@@ -32,6 +32,7 @@ use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::GenerationDebugLogContext;
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
use crate::http_server::endpoints::omni_api::generate::video::check_request::check_request;
use crate::http_server::endpoints::omni_api::generate::video::ingest_url_inputs::ingest_url_inputs;
@@ -217,10 +218,17 @@ pub async fn omni_api_video_generate_handler(
warn!("Failed to insert HTTP request debug log: {:?}", err);
}
// NB: Release the pooled DB connection before the (slow, external) generation call so we don't
// hold a pool slot idle while waiting on the provider — that's what starves the pool and causes
// PoolTimedOut on unrelated endpoints. We re-acquire below to write the result.
drop(mysql_connection);
// NB: The pipeline takes over the connection for its remaining pre-request DB writes (billing,
// outbound provider request debug log) and releases it before the (slow, external) generation
// call — holding a pool slot across that call is what starves the pool and causes PoolTimedOut
// on unrelated endpoints. We re-acquire below to write the result.
let debug_log_context = GenerationDebugLogContext {
event_token: &debug_log_event_token,
user_token,
ip_address: &ip_address,
request_url: &request_url,
};
let pipeline_result = run_pipeline_v2(RunPipelineV2Args {
router_builder: &router_builder,
@@ -229,6 +237,8 @@ pub async fn omni_api_video_generate_handler(
media_file_to_url_map: &media_file_to_url_map,
kinovi_character_id_map: &kinovi_character_id_map,
kinovi_account,
debug_log_context: &debug_log_context,
mysql_connection,
}).await;
// ==================== DEBUG LOG: PIPELINE ERROR ==================== //
@@ -258,45 +268,8 @@ pub async fn omni_api_video_generate_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
// ==================== DEBUG LOG: FAL REQUEST ==================== //
if let GenerateVideoResponse::Fal(ref fal_payload) = pipeline_result.response {
if let Some(ref outbound_request) = fal_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::FalRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Fal request debug log: {:?}", err);
}
}
}
// ==================== DEBUG LOG: GROK API REQUEST ==================== //
if let GenerateVideoResponse::Grok(ref grok_payload) = pipeline_result.response {
if let Some(ref outbound_request) = grok_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::GrokApiRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Grok API request debug log: {:?}", err);
}
}
}
// NB: Outbound provider requests (Fal/Grok/Kinovi) are debug-logged inside
// the pipeline BEFORE the send, so the payload is captured even on failure.
// ==================== WRITE RESULT ==================== //
@@ -15,6 +15,9 @@ use tokens::tokens::users::UserToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure;
use crate::http_server::endpoints::generate::common::generation_debug_logs::{
insert_provider_request_debug_log, provider_request_debug_log_type, GenerationDebugLogContext,
};
use crate::http_server::endpoints::omni_api::generate::video::helpers::bill_wallet::bill_wallet;
use crate::http_server::endpoints::omni_api::generate::video::helpers::build_router_client::build_router_client;
use crate::http_server::endpoints::omni_api::generate::video::helpers::pipeline_result::PipelineResult;
@@ -31,6 +34,11 @@ pub struct RunPipelineV2Args<'a> {
pub media_file_to_url_map: &'a Option<HashMap<MediaFileToken, String>>,
pub kinovi_character_id_map: &'a Option<HashMap<CharacterToken, String>>,
pub kinovi_account: KinoviAccount,
pub debug_log_context: &'a GenerationDebugLogContext<'a>,
/// The handler's open connection. The pipeline uses it for its remaining
/// pre-request DB writes (billing, outbound-request debug log) and releases
/// it BEFORE the external provider call.
pub mysql_connection: sqlx::pool::PoolConnection<sqlx::MySql>,
}
// NB: This pipeline does an external generation call (`upload_and_generate`) that can take many
@@ -45,6 +53,8 @@ pub async fn run_pipeline_v2(args: RunPipelineV2Args<'_>) -> Result<PipelineResu
media_file_to_url_map,
kinovi_character_id_map,
kinovi_account,
debug_log_context,
mut mysql_connection,
} = args;
let mut router_builder = router_builder.clone();
@@ -138,15 +148,25 @@ pub async fn run_pipeline_v2(args: RunPipelineV2Args<'_>) -> Result<PipelineResu
info!("v2 estimated cost: {} credits (estimates: {:?})", cost, cost_estimates);
// 3. Bill wallet (short-lived connection — released before the external call below).
let billing = {
let mut billing_connection = server_state.mysql_pool.acquire().await
.map_err(|err| {
error!("Failed to acquire MySQL connection for billing: {:?}", err);
CommonWebError::from_error(err)
})?;
bill_wallet(user_token, cost, &mut billing_connection).await?
};
// 3. Bill wallet on the handler's connection (same pre-request DB phase).
let billing = bill_wallet(user_token, cost, &mut mysql_connection).await?;
// Debug-log the outbound provider request BEFORE the send — still on the
// handler's connection — so the payload is captured even when the
// upload/enqueue fails.
if let Some(debug_log_type) = provider_request_debug_log_type(provider) {
insert_provider_request_debug_log(
debug_log_context,
debug_log_type,
&format!("{:#?}", draft_or_request),
&mut *mysql_connection,
).await;
}
// NB: Done with pre-request DB writes. Release the pooled connection before
// the (slow, external) provider call — holding it across that call is what
// starves the pool and causes PoolTimedOut. Post-send writes re-acquire.
drop(mysql_connection);
// 4. Upload media (if draft) and generate video.
// The entire block is wrapped so Kinovi failures trigger a refund.
@@ -27,6 +27,7 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::GenerationDebugLogContext;
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
use crate::http_server::endpoints::omni_gen::generate::image::hydrate_to_router_request::hydrate_to_router_request;
use crate::http_server::endpoints::omni_gen::generate::image::insert_db_job::insert_fal_job::{insert_fal_job, InsertFalJobArgs};
@@ -150,16 +151,25 @@ pub async fn omni_gen_image_generate_handler(
// ==================== PIPELINE ==================== //
// NB: Release the pooled DB connection before the (slow, external) generation call so we don't
// hold a pool slot idle while waiting on the provider — that's what starves the pool and causes
// PoolTimedOut on unrelated endpoints. We re-acquire below to write the result.
drop(mysql_connection);
// NB: The pipeline takes over the connection for its remaining pre-request DB writes (billing,
// outbound provider request debug log) and releases it before the (slow, external) generation
// call — holding a pool slot across that call is what starves the pool and causes PoolTimedOut
// on unrelated endpoints. We re-acquire below to write the result.
let debug_log_context = GenerationDebugLogContext {
event_token: &debug_log_event_token,
user_token,
ip_address: &ip_address,
request_url: &request_url,
};
let pipeline_result = run_pipeline_v2(RunPipelineV2Args {
router_builder: &router_builder,
server_state: &server_state,
user_token,
resolved_media: &resolved_media,
debug_log_context: &debug_log_context,
mysql_connection,
}).await;
// ==================== DEBUG LOG: PIPELINE ERROR ==================== //
@@ -189,25 +199,8 @@ pub async fn omni_gen_image_generate_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
// ==================== DEBUG LOG: FAL REQUEST ==================== //
if let GenerateImageResponse::Fal(ref fal_payload) = pipeline_result.response {
if let Some(ref outbound_request) = fal_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::FalRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Fal request debug log: {:?}", err);
}
}
}
// NB: Outbound provider requests (Fal/Grok/Kinovi) are debug-logged inside
// the pipeline BEFORE the send, so the payload is captured even on failure.
// ==================== WRITE RESULT ==================== //
@@ -17,6 +17,9 @@ use tokens::tokens::users::UserToken;
use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::{
insert_provider_request_debug_log, provider_request_debug_log_type, GenerationDebugLogContext,
};
use crate::http_server::endpoints::omni_gen::generate::image::pipeline_result::ImagePipelineResult;
use crate::http_server::endpoints::omni_gen::shared_utils::map_seedance2pro_router_error::map_router_error_to_web_error;
use crate::state::server_state::ServerState;
@@ -27,6 +30,11 @@ pub struct RunPipelineV2Args<'a> {
pub server_state: &'a ServerState,
pub user_token: &'a UserToken,
pub resolved_media: &'a MediaFilesAsCdnUrlListAndMap,
pub debug_log_context: &'a GenerationDebugLogContext<'a>,
/// The handler's open connection. The pipeline uses it for its remaining
/// pre-request DB writes (billing, outbound-request debug log) and releases
/// it BEFORE the external provider call.
pub mysql_connection: sqlx::pool::PoolConnection<sqlx::MySql>,
}
// NB: This pipeline does an external generation call (`finalize_and_generate`) that can take many
@@ -41,6 +49,8 @@ pub async fn run_pipeline_v2(
server_state,
user_token,
resolved_media,
debug_log_context,
mut mysql_connection,
} = args;
let hydrated_builder = apply_hydrated_media_inputs(
@@ -57,24 +67,36 @@ pub async fn run_pipeline_v2(
let apriori_job_token = InferenceJobToken::generate();
let maybe_wallet_ledger_entry_token = if cost > 0 {
// NB: Short-lived connection — released before the external call below.
let mut billing_connection = server_state.mysql_pool.acquire().await
.map_err(|err| {
warn!("Failed to acquire MySQL connection for image billing: {:?}", err);
CommonWebError::from_error(err)
})?;
// Billed on the handler's connection (same pre-request DB phase).
let deduction = attempt_wallet_deduction_else_common_web_error(
user_token,
Some(apriori_job_token.as_str()),
cost,
&mut billing_connection,
&mut mysql_connection,
).await?;
Some(deduction.ledger_entry_token)
} else {
None
};
// NB: No pooled DB connection is held across this external generation call.
// Debug-log the outbound provider request BEFORE the send — still on the
// handler's connection — so the payload is captured even when the
// upload/enqueue fails.
if let Some(debug_log_type) = provider_request_debug_log_type(draft_or_request.get_provider()) {
insert_provider_request_debug_log(
debug_log_context,
debug_log_type,
&format!("{:#?}", draft_or_request),
&mut *mysql_connection,
).await;
}
// NB: Done with pre-request DB writes. Release the pooled connection before
// the (slow, external) provider call — holding it across that call is what
// starves the pool and causes PoolTimedOut. No pooled DB connection is held
// across the external generation call below.
drop(mysql_connection);
let response = finalize_and_generate(draft_or_request, server_state, resolved_media).await?;
Ok(ImagePipelineResult {
@@ -31,6 +31,7 @@ use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoints::generate::common::generation_debug_logs::GenerationDebugLogContext;
use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
use crate::http_server::endpoints::omni_gen::generate::video::helpers::hydrate_router_request::hydrate_to_router_request;
use crate::http_server::endpoints::omni_gen::generate::video::helpers::resolve_kinovi_character_ids::resolve_kinovi_character_ids;
@@ -209,10 +210,17 @@ pub async fn omni_gen_video_generate_handler(
warn!("Failed to insert HTTP request debug log: {:?}", err);
}
// NB: Release the pooled DB connection before the (slow, external) generation call so we don't
// hold a pool slot idle while waiting on the provider — that's what starves the pool and causes
// PoolTimedOut on unrelated endpoints. We re-acquire below to write the result.
drop(mysql_connection);
// NB: The pipeline takes over the connection for its remaining pre-request DB writes (billing,
// outbound provider request debug log) and releases it before the (slow, external) generation
// call — holding a pool slot across that call is what starves the pool and causes PoolTimedOut
// on unrelated endpoints. We re-acquire below to write the result.
let debug_log_context = GenerationDebugLogContext {
event_token: &debug_log_event_token,
user_token,
ip_address: &ip_address,
request_url: &request_url,
};
let pipeline_result = run_pipeline_v2(RunPipelineV2Args {
router_builder: &router_builder,
@@ -221,6 +229,8 @@ pub async fn omni_gen_video_generate_handler(
media_file_to_url_map: &media_file_to_url_map,
kinovi_character_id_map: &kinovi_character_id_map,
kinovi_account,
debug_log_context: &debug_log_context,
mysql_connection,
}).await;
// ==================== DEBUG LOG: PIPELINE ERROR ==================== //
@@ -250,45 +260,8 @@ pub async fn omni_gen_video_generate_handler(
let mut mysql_connection = server_state.mysql_pool.acquire().await?;
// ==================== DEBUG LOG: FAL REQUEST ==================== //
if let GenerateVideoResponse::Fal(ref fal_payload) = pipeline_result.response {
if let Some(ref outbound_request) = fal_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::FalRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Fal request debug log: {:?}", err);
}
}
}
// ==================== DEBUG LOG: GROK API REQUEST ==================== //
if let GenerateVideoResponse::Grok(ref grok_payload) = pipeline_result.response {
if let Some(ref outbound_request) = grok_payload.maybe_outbound_request {
if let Err(err) = insert_debug_log(InsertDebugLogArgs {
apriori_debug_log_event_token: Some(&debug_log_event_token),
maybe_creator_user_token: Some(user_token),
debug_log_type: DebugLogType::GrokApiRequest,
maybe_log_level: Some(DebugLogLevel::Info),
maybe_ip_address: Some(&ip_address),
maybe_url: Some(&request_url),
message: &format!("{:#?}", outbound_request),
mysql_executor: &mut *mysql_connection,
phantom: Default::default(),
}).await {
warn!("Failed to insert Grok API request debug log: {:?}", err);
}
}
}
// NB: Outbound provider requests (Fal/Grok/Kinovi) are debug-logged inside
// the pipeline BEFORE the send, so the payload is captured even on failure.
// ==================== WRITE RESULT ==================== //
@@ -15,6 +15,9 @@ use tokens::tokens::users::UserToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure;
use crate::http_server::endpoints::generate::common::generation_debug_logs::{
insert_provider_request_debug_log, provider_request_debug_log_type, GenerationDebugLogContext,
};
use crate::http_server::endpoints::omni_gen::generate::video::helpers::bill_wallet::bill_wallet;
use crate::http_server::endpoints::omni_gen::generate::video::helpers::build_router_client::build_router_client;
use crate::http_server::endpoints::omni_gen::generate::video::helpers::pipeline_result::PipelineResult;
@@ -31,6 +34,11 @@ pub struct RunPipelineV2Args<'a> {
pub media_file_to_url_map: &'a Option<HashMap<MediaFileToken, String>>,
pub kinovi_character_id_map: &'a Option<HashMap<CharacterToken, String>>,
pub kinovi_account: KinoviAccount,
pub debug_log_context: &'a GenerationDebugLogContext<'a>,
/// The handler's open connection. The pipeline uses it for its remaining
/// pre-request DB writes (billing, outbound-request debug log) and releases
/// it BEFORE the external provider call.
pub mysql_connection: sqlx::pool::PoolConnection<sqlx::MySql>,
}
// NB: This pipeline does an external generation call (`upload_and_generate`) that can take many
@@ -45,6 +53,8 @@ pub async fn run_pipeline_v2(args: RunPipelineV2Args<'_>) -> Result<PipelineResu
media_file_to_url_map,
kinovi_character_id_map,
kinovi_account,
debug_log_context,
mut mysql_connection,
} = args;
let mut router_builder = router_builder.clone();
@@ -141,15 +151,25 @@ pub async fn run_pipeline_v2(args: RunPipelineV2Args<'_>) -> Result<PipelineResu
info!("v2 estimated cost: {} credits (estimates: {:?})", cost, cost_estimates);
// 3. Bill wallet (short-lived connection — released before the external call below).
let billing = {
let mut billing_connection = server_state.mysql_pool.acquire().await
.map_err(|err| {
error!("Failed to acquire MySQL connection for billing: {:?}", err);
CommonWebError::from_error(err)
})?;
bill_wallet(user_token, cost, &mut billing_connection).await?
};
// 3. Bill wallet on the handler's connection (same pre-request DB phase).
let billing = bill_wallet(user_token, cost, &mut mysql_connection).await?;
// Debug-log the outbound provider request BEFORE the send — still on the
// handler's connection — so the payload is captured even when the
// upload/enqueue fails.
if let Some(debug_log_type) = provider_request_debug_log_type(provider) {
insert_provider_request_debug_log(
debug_log_context,
debug_log_type,
&format!("{:#?}", draft_or_request),
&mut *mysql_connection,
).await;
}
// NB: Done with pre-request DB writes. Release the pooled connection before
// the (slow, external) provider call — holding it across that call is what
// starves the pool and causes PoolTimedOut. Post-send writes re-acquire.
drop(mysql_connection);
// 4. Upload media (if draft) and generate video.
// The entire block is wrapped so Kinovi failures trigger a refund.
@@ -1,15 +1,15 @@
use std::sync::Arc;
use actix_web::{http::StatusCode, web, HttpRequest, HttpResponse, ResponseError};
use chrono::NaiveDateTime;
use log::{debug, error, warn};
use mysql_queries::queries::stats::get_unified_queue_stats::get_unified_queue_stats;
use redis_common::redis_cache_keys::RedisCacheKeys;
use actix_web::web::Json;
use chrono::{NaiveDateTime, Utc};
use utoipa::ToSchema;
use crate::http_server::endpoints::stats::result_transformer::{database_result_to_cacheable, CacheableQueueStats};
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::state::server_state::ServerState;
/// The queue counts are frozen at the values from when the legacy inference
/// systems were shut down. Legacy clients still poll this endpoint, so it
/// stays mounted, but it no longer consults the database.
const FROZEN_TOTAL_PENDING_JOB_COUNT: u64 = 2_409_539;
const FROZEN_PENDING_TACOTRON2_JOBS: u64 = 1_037_783;
const FROZEN_PENDING_RVC_JOBS: u64 = 288_922;
const FROZEN_PENDING_SVC_JOBS: u64 = 27_417;
const REFRESH_INTERVAL_MILLIS: u64 = 15_000;
#[derive(Serialize, ToSchema)]
pub struct GetUnifiedQueueStatsSuccessResponse {
@@ -17,7 +17,6 @@ pub struct GetUnifiedQueueStatsSuccessResponse {
pub cache_time: NaiveDateTime,
/// Tell the frontend client how fast to refresh their view of this list.
/// During an attack, we may want this to go extremely slow.
pub refresh_interval_millis: u64,
pub inference: ModernInferenceQueueStats,
@@ -33,7 +32,6 @@ pub struct LegacyQueueDetails {
pub struct ModernInferenceQueueStats {
pub total_pending_job_count: u64,
#[deprecated(note="the frontend uses this field, but we should switch to total_pending_job_count")]
pub pending_job_count: u64,
pub by_queue: ByQueueStats,
@@ -58,200 +56,40 @@ pub struct ByQueueStats {
pub pending_acting_face: u64,
}
#[derive(Debug, ToSchema)]
pub enum GetUnifiedQueueStatsError {
ServerError,
}
impl ResponseError for GetUnifiedQueueStatsError {
fn status_code(&self) -> StatusCode {
match *self {
GetUnifiedQueueStatsError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR,
}
}
fn error_response(&self) -> HttpResponse {
let error_reason = match self {
GetUnifiedQueueStatsError::ServerError => "server error".to_string(),
};
to_simple_json_error(&error_reason, self.status_code())
}
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
impl std::fmt::Display for GetUnifiedQueueStatsError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?}", self)
}
}
/// Get queue stats for most inference jobs (tts, voice conversion, comfy, etc.)
/// [DEPRECATED] Get queue stats for legacy inference jobs (tts, voice conversion, etc.)
///
/// The legacy inference systems are gone; this returns a hardcoded snapshot
/// for old clients that still poll it.
#[utoipa::path(
get,
tag = "Stats",
path = "/v1/stats/queues",
responses(
(status = 200, description = "Success", body = GetUnifiedQueueStatsSuccessResponse),
(status = 500, description = "Server error", body = GetUnifiedQueueStatsError),
),
)]
pub async fn get_unified_queue_stats_handler(
http_request: HttpRequest,
server_state: web::Data<Arc<ServerState>>
) -> Result<HttpResponse, GetUnifiedQueueStatsError> {
if server_state.flags.disable_unified_queue_stats_endpoint {
// NB: Despite the cache being a powerful protector of the database (this is an expensive query),
// if the cache goes stale during an outage, there is no protection. This feature flag lets us
// shut off all traffic to the endpoint.
return render_response_busy(GetUnifiedQueueStatsSuccessResponse {
success: true,
cache_time: NaiveDateTime::from_timestamp(0, 0),
refresh_interval_millis: server_state.flags.frontend_unified_queue_stats_refresh_interval_millis,
inference: ModernInferenceQueueStats {
total_pending_job_count: 10_000,
pending_job_count: 10_000,
by_queue: ByQueueStats {
pending_face_animation_jobs: 10_000,
pending_storyteller_studio: 10_000,
pending_rvc_jobs: 10_000,
pending_svc_jobs: 10_000,
pending_tacotron2_jobs: 10_000,
pending_voice_designer: 10_000,
pending_stable_diffusion: 10_000,
pending_acting_face: 10_000,
}
},
legacy_tts: LegacyQueueDetails {
pending_job_count: 10_000,
},
});
}
let maybe_cached = server_state.caches.ephemeral.queue_stats.grab_copy_without_bump_if_unexpired()
.map_err(|e| {
error!("error consulting cache: {:?}", e);
GetUnifiedQueueStatsError::ServerError
})?;
let cacheable_stats_result = match maybe_cached {
Some(cached) => {
debug!("serving from in-memory cache");
cached
},
None => {
debug!("populating unified queue stats from Redis *OR* database");
let mysql_pool = server_state.mysql_pool.clone();
let get_stats = move || {
// NB: async closures are not yet stable in Rust, so we include an async block.
async move {
debug!("querying from database...");
let db_result = get_unified_queue_stats(
&mysql_pool
).await;
db_result
.map(|db_result| database_result_to_cacheable(db_result))
.map(|cacheable_result| Some(cacheable_result)) // TODO/FIXME: Make a better redis cache
}
};
// NB(2023-07-27): Double layers of caching (in-memory + Redis) is probably exorbitant, but this seems fine for now.
// This endpoint's query (even when in-memory cached across the cluster) was causing the DB CPU to hit 100%.
let stats_query_result = match server_state.redis_ttl_cache.get_connection() {
Err(err) => {
warn!("Error loading Redis connection from TTL cache (calling DB instead): {:?}", err);
get_stats().await
}
Ok(mut redis_ttl_cache) => {
let cache_key = RedisCacheKeys::get_unified_queue_stats_endpoint();
redis_ttl_cache.lazy_load_if_not_cached(&cache_key, move || {
get_stats()
}).await
}
};
match stats_query_result {
// If the database misbehaves (eg. DDoS), let's stop spamming it.
// We'll attempt to read the old value from the cache and keep going.
Err(err) => {
warn!("error querying database / inserting into cache: {:?}", err);
let maybe_cached = server_state.caches.ephemeral.queue_stats.grab_even_expired_and_bump()
.map_err(|err| {
error!("error consulting cache (even expired): {:?}", err);
GetUnifiedQueueStatsError::ServerError
})?;
maybe_cached.ok_or_else(|| {
error!("error querying database and subsequently reading cache: {:?}", err);
GetUnifiedQueueStatsError::ServerError
})?
}
// Happy path...
Ok(Some(cacheable_stats)) => {
server_state.caches.ephemeral.queue_stats.store_copy(&cacheable_stats)
.map_err(|e| {
error!("error storing cache: {:?}", e);
GetUnifiedQueueStatsError::ServerError
})?;
cacheable_stats
}
// TODO/FIXME(bt,2023-07-30): This match arm should never happen.
Ok(None) => CacheableQueueStats::default()
}
},
};
render_response_ok(GetUnifiedQueueStatsSuccessResponse {
#[deprecated(note = "legacy inference is shut down; this endpoint returns a hardcoded snapshot")]
pub async fn get_unified_queue_stats_handler() -> Json<GetUnifiedQueueStatsSuccessResponse> {
Json(GetUnifiedQueueStatsSuccessResponse {
success: true,
cache_time: cacheable_stats_result.cache_time,
refresh_interval_millis: server_state.flags.frontend_pending_inference_refresh_interval_millis,
cache_time: Utc::now().naive_utc(),
refresh_interval_millis: REFRESH_INTERVAL_MILLIS,
inference: ModernInferenceQueueStats {
total_pending_job_count: cacheable_stats_result.queues.total_generic,
pending_job_count: cacheable_stats_result.queues.total_generic,
total_pending_job_count: FROZEN_TOTAL_PENDING_JOB_COUNT,
pending_job_count: FROZEN_TOTAL_PENDING_JOB_COUNT,
by_queue: ByQueueStats {
pending_face_animation_jobs: cacheable_stats_result.queues.sad_talker,
pending_rvc_jobs: cacheable_stats_result.queues.rvc_v2,
pending_svc_jobs: cacheable_stats_result.queues.so_vits_svc,
pending_tacotron2_jobs: cacheable_stats_result.queues.tacotron2,
pending_voice_designer: cacheable_stats_result.queues.vall_e_x,
pending_stable_diffusion: cacheable_stats_result.queues.stable_diffusion,
pending_storyteller_studio: cacheable_stats_result.queues.storyteller_studio,
pending_acting_face: cacheable_stats_result.queues.acting_face,
}
pending_tacotron2_jobs: FROZEN_PENDING_TACOTRON2_JOBS,
pending_voice_designer: 0,
pending_rvc_jobs: FROZEN_PENDING_RVC_JOBS,
pending_svc_jobs: FROZEN_PENDING_SVC_JOBS,
pending_stable_diffusion: 0,
pending_face_animation_jobs: 0,
pending_storyteller_studio: 0,
pending_acting_face: 0,
},
},
legacy_tts: LegacyQueueDetails {
pending_job_count: cacheable_stats_result.queues.legacy_tts,
pending_job_count: 0,
},
})
}
pub fn render_response_busy(response: GetUnifiedQueueStatsSuccessResponse) -> Result<HttpResponse, GetUnifiedQueueStatsError> {
let body = render_response_payload(response)?;
Ok(HttpResponse::TooManyRequests()
.content_type("application/json")
.body(body))
}
pub fn render_response_ok(response: GetUnifiedQueueStatsSuccessResponse) -> Result<HttpResponse, GetUnifiedQueueStatsError> {
let body = render_response_payload(response)?;
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(body))
}
pub fn render_response_payload(response: GetUnifiedQueueStatsSuccessResponse) -> Result<String, GetUnifiedQueueStatsError> {
let body = serde_json::to_string(&response)
.map_err(|e| {
error!("error returning response: {:?}", e);
GetUnifiedQueueStatsError::ServerError
})?;
Ok(body)
}
@@ -1,2 +1 @@
pub mod get_unified_queue_stats_handler;
pub mod result_transformer;
@@ -1,82 +0,0 @@
use std::collections::HashMap;
use chrono::NaiveDateTime;
use mysql_queries::queries::stats::get_unified_queue_stats::QueueStatsRow;
#[derive(Serialize, Deserialize, Clone)]
pub struct CacheableQueueStats {
pub cache_time: NaiveDateTime,
pub queues: Queues,
}
impl Default for CacheableQueueStats {
fn default() -> Self {
Self {
cache_time: NaiveDateTime::from_timestamp(0, 0),
queues: Default::default(),
}
}
}
#[derive(Serialize, Deserialize, Clone, Default)]
pub struct Queues {
// Text to Speech
pub legacy_tts: u64, // Tacotron2 (as tts_inference_jobs)
pub total_generic: u64, // Everything except tts_inference_jobs' Tacotron2 summed together
pub tacotron2: u64, // Tacotron2 (as generic_inference_jobs; NB: included in `total_generic`)
pub vall_e_x: u64, // Counts both model weight calculation and inference.
// Voice Conversion
pub rvc_v2: u64,
pub so_vits_svc: u64,
// Image
pub stable_diffusion: u64,
// Video
pub sad_talker: u64,
pub storyteller_studio: u64,
pub acting_face: u64,
}
pub fn database_result_to_cacheable(database_records: Vec<QueueStatsRow>) -> CacheableQueueStats {
let database_time = database_records.first()
.map(|record| record.present_time)
.unwrap_or(NaiveDateTime::from_timestamp(0, 0));
// Queue lengths
let mut queue_lengths : HashMap<String, u64> = HashMap::new();
for record in database_records.into_iter() {
queue_lengths.insert(record.queue_type, record.pending_job_count);
}
CacheableQueueStats {
cache_time: database_time,
queues: Queues {
// Statistic rollup.
total_generic: queue_lengths.iter()
.filter(|(queue, _count)| !(*queue).eq("legacy_tts"))
.map(|(_queue, count)| count)
.sum(),
// TTS
legacy_tts: unwrap(queue_lengths.get("legacy_tts")),
tacotron2: unwrap(queue_lengths.get("tacotron2")),
vall_e_x: unwrap(queue_lengths.get("vall_e_x")),
// VC
rvc_v2: unwrap(queue_lengths.get("rvc_v2")),
so_vits_svc: unwrap(queue_lengths.get("so_vits_svc")),
// Image
stable_diffusion: unwrap(queue_lengths.get("stable_diffusion")),
// Video
sad_talker: unwrap(queue_lengths.get("sad_talker")),
storyteller_studio: unwrap(queue_lengths.get("video_render")),
acting_face: unwrap(queue_lengths.get("live_portrait")),
},
}
}
fn unwrap(inner: Option<&u64>) -> u64 {
inner.map(|i|*i).unwrap_or(0)
}
@@ -1,146 +0,0 @@
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::web::Json;
use actix_web::http::StatusCode;
use actix_web::web::Path;
use actix_web::{web, HttpRequest};
use log::warn;
use artcraft_api_defs::common::responses::simple_generic_json_success::SimpleGenericJsonSuccess;
use mysql_queries::queries::tts::tts_results::delete_tts_result_various_scopes::delete_tts_inference_result_as_mod;
use mysql_queries::queries::tts::tts_results::delete_tts_result_various_scopes::delete_tts_inference_result_as_user;
use mysql_queries::queries::tts::tts_results::delete_tts_result_various_scopes::undelete_tts_inference_result_as_mod;
use mysql_queries::queries::tts::tts_results::delete_tts_result_various_scopes::undelete_tts_inference_result_as_user;
use mysql_queries::queries::tts::tts_results::query_tts_result::select_tts_result_by_token;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::state::server_state::ServerState;
use crate::util::delete_role_disambiguation::delete_role_disambiguation;
use crate::util::delete_role_disambiguation::DeleteRole;
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct DeleteTtsInferenceResultPathInfo {
token: String,
}
#[derive(Deserialize)]
pub struct DeleteTtsInferenceResultRequest {
set_delete: bool,
/// NB: this is only to disambiguate when a user is both a mod and an author.
as_mod: Option<bool>,
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
pub async fn delete_tts_inference_result_handler(
http_request: HttpRequest,
path: Path<DeleteTtsInferenceResultPathInfo>,
request: Json<DeleteTtsInferenceResultRequest>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<SimpleGenericJsonSuccess>, CommonWebError> {
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session(&http_request, &server_state.mysql_pool)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from_error(e)
})?;
let user_session = match maybe_user_session {
Some(session) => session,
None => {
warn!("not logged in");
return Err(CommonWebError::NotAuthorized);
}
};
// NB: First permission check.
// Only mods should see deleted models (both user_* and mod_* deleted).
let is_mod_that_can_see_deleted = user_session.can_delete_other_users_tts_results;
let inference_result_query_result = select_tts_result_by_token(
&path.token,
is_mod_that_can_see_deleted,
&server_state.mysql_pool,
).await;
let tts_inference_result = match inference_result_query_result {
Err(e) => {
warn!("query error: {:?}", e);
return Err(CommonWebError::from_anyhow_error(e));
}
Ok(None) => return Err(CommonWebError::NotFound),
Ok(Some(inference_result)) => inference_result,
};
// NB: Second set of permission checks
let is_author = tts_inference_result.maybe_creator_user_token
.as_ref()
.map(|creator_token| creator_token == &user_session.user_token)
.unwrap_or(false);
let is_mod = user_session.can_delete_other_users_tts_results;
if !is_author && !is_mod {
warn!("user is not allowed to delete inference results: {:?}", user_session.user_token);
return Err(CommonWebError::NotAuthorized);
}
let delete_role = delete_role_disambiguation(is_mod, is_author, request.as_mod);
let query_result = if request.set_delete {
match delete_role {
DeleteRole::ErrorDoNotDelete => {
warn!("user is not allowed to delete inference results: {:?}", user_session.user_token);
return Err(CommonWebError::NotAuthorized);
}
DeleteRole::AsUser => {
delete_tts_inference_result_as_user(
&path.token,
&server_state.mysql_pool
).await
}
DeleteRole::AsMod => {
delete_tts_inference_result_as_mod(
&path.token,
user_session.user_token.as_str(),
&server_state.mysql_pool
).await
}
}
} else {
match delete_role {
DeleteRole::ErrorDoNotDelete => {
warn!("user is not allowed to undelete inference results: {:?}", user_session.user_token);
return Err(CommonWebError::NotAuthorized);
}
DeleteRole::AsUser => {
// NB: Technically only mods can see their own inference_results here
undelete_tts_inference_result_as_user(
&path.token,
&server_state.mysql_pool
).await
}
DeleteRole::AsMod => {
undelete_tts_inference_result_as_mod(
&path.token,
user_session.user_token.as_str(),
&server_state.mysql_pool
).await
}
}
};
match query_result {
Ok(_) => {},
Err(err) => {
warn!("Update tts mod approval status DB error: {:?}", err);
return Err(CommonWebError::from_anyhow_error(err));
}
};
Ok(Json(SimpleGenericJsonSuccess { success: true }))
}
@@ -1,86 +1,14 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::web::Json;
use actix_web::{web, HttpRequest, HttpResponse};
use log::{info, warn};
use rand::seq::IndexedRandom;
use rand::Rng;
use redis::Commands;
use utoipa::ToSchema;
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
use enums::by_table::generic_inference_jobs::inference_job_product_category::InferenceJobProductCategory;
use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
use enums::common::visibility::Visibility;
use http_server_common::request::get_request_api_token::get_request_api_token;
use http_server_common::request::get_request_header_optional::get_request_header_optional;
use http_server_common::request::get_request_ip::get_request_ip;
use migration::text_to_speech::get_tts_model_for_enqueue_inference_migration::TtsModelForEnqueueInferenceMigrationWrapper;
use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs};
use mysql_queries::payloads::generic_inference_args::inner_payloads::gptsovits_payload::GptSovitsPayload;
use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs};
use redis_common::redis_keys::RedisKeys;
use tokens::tokens::users::UserToken;
use tts_common::priority::{FAKEYOU_DEFAULT_VALID_API_TOKEN_PRIORITY_LEVEL, FAKEYOU_INVESTOR_PRIORITY_LEVEL};
use user_input_common::check_for_slurs::contains_slurs;
use crate::configs::app_startup::username_set::UsernameSet;
use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session;
use crate::configs::plans::plan::Plan;
use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::get_tts_model_with_caching::get_tts_model_with_caching;
use crate::http_server::user_lookup::user_session::session_utils::lookup::user_session_extended::UserSessionExtended;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::state::server_state::ServerState;
// TODO: Temporary for investor demo
const STORYTELLER_DEMO_COOKIE_NAME : &str = "storyteller_demo";
/// Debug requests mean different things, depending on the type of worker (the legacy
/// tts-inference-job vs. the new inference-job):
///
/// * Legacy (tts): these requests get routed to special "debug-only" workers, which can
/// be used to trial new code, run debugging, etc.
///
/// * New (generic): the flag gets associated with the request which may be used for
/// anything in the code. (There's a new system for routing requests called
/// "routing tags" that is independent of this flag.)
///
const DEBUG_HEADER_NAME : &str = "enable-debug-mode";
/// The routing tag header can send workloads to particular k8s hosts.
/// This is useful for catching the live logs or intercepting the job.
/// NB: This is only for the new job system.
const ROUTING_TAG_HEADER_NAME : &str = "routing-tag";
const USER_FAKEYOU_USER_TOKEN : &str = "U:N5J8JXPW9BTYX";
const USER_NEWS_STORY_USER_TOKEN : &str = "U:XAWRARC1N89X6";
/// Safety guard against plans supporting too lengthy TTS inference text.
const MAX_TTS_LENGTH : usize = 10_000;
#[derive(Deserialize, ToSchema)]
pub struct InferTtsRequest {
uuid_idempotency_token: String,
tts_model_token: String,
inference_text: String,
creator_set_visibility: Option<Visibility>,
is_storyteller_demo: Option<bool>,
}
/// Sentinel job token returned to legacy clients now that TTS inference is shut down.
/// There is no real job behind it; status polls for it will never complete.
const SYNTHETIC_INFERENCE_JOB_TOKEN: &str = "synthetic_too_many_requests";
#[derive(Serialize, ToSchema)]
pub struct InferTtsSuccessResponse {
pub success: bool,
pub inference_job_token: String,
// This is the type of inference job token returned.
pub inference_job_token_type: InferenceJobTokenType,
}
@@ -91,537 +19,27 @@ pub enum InferenceJobTokenType {
#[serde(rename = "legacy_tts")]
LegacyTts,
/// Modern shared inference type (everything will use this in the
/// future, and this endpoint will die.)
/// Modern shared inference type
#[serde(rename = "generic")]
Generic,
}
#[derive(Debug, ToSchema)]
pub enum InferTtsError {
BadInput(String),
NotAuthorized,
ServerError,
RateLimited,
NotFound,
}
impl ResponseError for InferTtsError {
fn status_code(&self) -> StatusCode {
match *self {
InferTtsError::BadInput(_) => StatusCode::BAD_REQUEST,
InferTtsError::NotAuthorized => StatusCode::UNAUTHORIZED,
InferTtsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
InferTtsError::RateLimited => StatusCode::TOO_MANY_REQUESTS,
InferTtsError::NotFound => StatusCode::NOT_FOUND,
}
}
fn error_response(&self) -> HttpResponse {
let error_reason = match self {
InferTtsError::BadInput(reason) => reason.to_string(),
InferTtsError::NotAuthorized => "unauthorized".to_string(),
InferTtsError::ServerError => "server error".to_string(),
InferTtsError::RateLimited => "rate limited".to_string(),
InferTtsError::NotFound => "not found".to_string(),
};
to_simple_json_error(&error_reason, self.status_code())
}
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
impl fmt::Display for InferTtsError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self)
}
}
/// [DEPRECATED] Enqueue a legacy TTS inference request.
///
/// Legacy TTS inference is shut down. This endpoint unconditionally returns a
/// canned success response with a synthetic job token so old clients don't error.
#[utoipa::path(
post,
tag = "TTS",
path = "/v1/tts/inference",
responses(
(status = 200, description = "Success response", body = InferTtsSuccessResponse),
(status = 400, description = "Bad input", body = InferTtsError),
(status = 401, description = "Not authorized", body = InferTtsError),
(status = 429, description = "Rate limited", body = InferTtsError),
(status = 500, description = "Server error", body = InferTtsError)
),
params(("request" = InferTtsRequest, description = "Payload for Request"))
)]
pub async fn enqueue_infer_tts_handler(
http_request: HttpRequest,
request: Json<InferTtsRequest>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<InferTtsSuccessResponse>, InferTtsError> {
if server_state.flags.disable_tts {
return Err(InferTtsError::RateLimited);
}
let mut is_from_api = false;
let mut maybe_user_token : Option<String> = None;
let mut priority_level ;
let mut use_high_priority_rate_limiter = false; // NB: Careful!
let mut disable_rate_limiter = false; // NB: Careful!
let mut mysql_connection = server_state.mysql_pool
.acquire()
.await
.map_err(|err| {
warn!("MySql pool error: {:?}", err);
InferTtsError::ServerError
})?;
// ==================== USER SESSION ==================== //
let maybe_user_session : Option<UserSessionExtended> = server_state
.session_checker
.maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
InferTtsError::ServerError
})?;
if let Some(user_session) = maybe_user_session.as_ref() {
maybe_user_token = Some(user_session.user_token.to_string());
}
// TODO: Plan should handle "first anonymous use" and "investor" cases.
let plan = get_correct_plan_for_session(
server_state.server_environment,
maybe_user_session.as_ref());
priority_level = plan.tts_base_priority_level();
// ==================== API TOKENS ==================== //
if let Some(api_token) = get_request_api_token(&http_request) {
let maybe_api_token_configs = server_state
.static_api_token_set
.get_api_token(&api_token);
if let Some(api_token_configs) = maybe_api_token_configs {
is_from_api = true;
priority_level = api_token_configs.maybe_priority_level
.unwrap_or(FAKEYOU_DEFAULT_VALID_API_TOKEN_PRIORITY_LEVEL);
use_high_priority_rate_limiter = api_token_configs.maybe_use_high_priority_rate_limiter
.unwrap_or(false);
disable_rate_limiter = api_token_configs.maybe_disable_rate_limiter.unwrap_or(false);
if let Some(user_token_override) = api_token_configs.maybe_user_token.as_deref() {
maybe_user_token = Some(user_token_override.trim().to_string());
}
}
}
// ==================== INVESTOR PRIORITY ==================== //
// TODO/TEMP: Give investors even more priority
let mut is_investor = false;
// TODO/TEMP: The storyteller.io website's AJAX calls will set this.
// This is just for the YCombinator demo.
match request.is_storyteller_demo {
Some(true) => {
is_investor = true;
},
_ => {},
};
if is_investor {
priority_level = FAKEYOU_INVESTOR_PRIORITY_LEVEL;
}
// ==================== ENGINEERING FLAGS / HEADERS: DEBUG MODE, ETC. ==================== //
let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME)
.is_some();
let maybe_routing_tag=
get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME)
.map(|routing_tag| routing_tag.trim().to_string());
// ==================== BANNED USERS ==================== //
if let Some(ref user) = maybe_user_session {
if user.role.is_banned {
warn!("User is not authorized to use TTS model because they are banned.");
return Err(InferTtsError::NotAuthorized);
}
}
// ==================== RATE LIMIT ==================== //
if !disable_rate_limiter {
let maybe_username = maybe_user_session
.as_ref()
.map(|session| session.user.username.as_str());
// FIXME(bt,2023-12-13): These boolean flags are a bit dangerous.
let rate_limiter_type = get_rate_limiter_type(
maybe_username,
&server_state.redis_rate_limiters.api_ai_streamer_username_set,
is_investor,
is_from_api,
use_high_priority_rate_limiter
);
let rate_limiter = match rate_limiter_type {
RateLimiterType::LoggedOut => &server_state.redis_rate_limiters.logged_out,
RateLimiterType::LoggedIn => &server_state.redis_rate_limiters.logged_in,
RateLimiterType::ApiHighPriority => {
info!("Using API high priority rate limiter");
&server_state.redis_rate_limiters.api_high_priority
},
RateLimiterType::ApiAiStreamer => {
info!("Using AI streamer rate limiter");
&server_state.redis_rate_limiters.api_ai_streamers
},
};
if let Err(_err) = rate_limiter.rate_limit_request(&http_request).await {
return Err(InferTtsError::RateLimited);
}
}
// ==================== CHECK TTS MODEL EXISTENCE / SETTINGS==================== //
let tts_model_lookup_result = get_tts_model_with_caching(
&request.tts_model_token,
&server_state.redis_ttl_cache,
&mut mysql_connection,
).await;
let tts_model = match tts_model_lookup_result {
Ok(Some(model)) => model,
Ok(None) => {
warn!("Model could not be found: {}", request.tts_model_token);
return Err(InferTtsError::NotFound);
}
Err(err) => {
warn!("Error loading TTS model: {:?}", err);
return Err(InferTtsError::ServerError);
}
};
let is_authorized = check_if_authorized_to_use_model(
maybe_user_session.as_ref(),
&tts_model,
).await;
if !is_authorized {
warn!("User is not authorized to use TTS model due to checks.");
return Err(InferTtsError::NotAuthorized);
}
// ==================== CHECK INFERENCE TEXT ==================== //
let ip_address = get_request_ip(&http_request);
let mut inference_text = request.inference_text.trim().to_string();
if let Err(reason) = validate_inference_text(&inference_text, &plan) {
return Err(InferTtsError::BadInput(reason));
}
if contains_slurs(&inference_text) {
return Err(InferTtsError::BadInput("text contains slurs".to_string()));
}
if let Some(user) = maybe_user_session.as_ref() {
let is_user_token_troll_banned = server_state.troll_bans.user_tokens
.contains_user_token(&user.user_token)
.unwrap_or(false); // NB: Fail open
let is_ip_address_troll_banned = server_state.troll_bans.ip_addresses
.contains_ip_address(&ip_address)
.unwrap_or(false); // NB: Fail open
let is_troll_banned = is_user_token_troll_banned || is_ip_address_troll_banned;
if is_troll_banned {
let test = server_state.flags.troll_ban_user_percent as i32; // NB: Should be 0 - 100.
let num = rand::thread_rng().gen_range(0..100);
if test > num {
warn!("Generating troll response for user token: {}", user.user_token);
inference_text = random_troll_text().to_string();
}
}
}
// ==================== CHECK AND PREPARE TTS ==================== //
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
warn!("redis error: {:?}", e);
InferTtsError::ServerError
})?;
let redis_count_key = RedisKeys::tts_model_usage_count(&request.tts_model_token);
redis.incr::<_, _, ()>(&redis_count_key, 1)
.map_err(|e| {
warn!("redis error: {:?}", e);
InferTtsError::ServerError
})?;
let maybe_user_preferred_visibility : Option<Visibility> = maybe_user_session
.as_ref()
.map(|user_session| user_session.preferences.preferred_tts_result_visibility);
let set_visibility = request.creator_set_visibility
.or(maybe_user_preferred_visibility)
.unwrap_or(Visibility::Public);
// ==================== ENQUEUE APPROPRIATE TTS TYPE ==================== //
let max_duration_seconds = plan.tts_max_duration_seconds();
// This branch uses the `generic-inference-job` service and tables.
info!("Creating tts inference job record (new generic job system)...");
let maybe_creator_user_token_typed = maybe_user_token
.as_deref()
.map(|token| UserToken::new_from_str(token));
let maybe_avt_token = server_state
.avt_cookie_manager
.get_avt_token_from_request(&http_request);
let maybe_model_type = match tts_model.job_type() {
InferenceJobType::GptSovits => None, // TODO(bt,2024-08-07): why don't we set this?
_ => Some(InferenceModelType::Tacotron2)
};
let maybe_inference_args = match tts_model.job_type() {
InferenceJobType::Tacotron2 => Some(GenericInferenceArgs {
inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech),
args: None, // NB: We don't need to encode any args yet for TT2.
}),
InferenceJobType::GptSovits => Some(GenericInferenceArgs {
inference_category: Some(InferenceCategoryAbbreviated::GptSovits),
args: Some(PolymorphicInferenceArgs::Gs(GptSovitsPayload {
append_advertisement: None,
// Model download only args:
creator_visibility: None,
maybe_title: None,
maybe_description: None,
})),
}),
_ => None, // NB: Shouldn't occur
};
let maybe_product_category = match tts_model.job_type() {
InferenceJobType::Tacotron2 => Some(InferenceJobProductCategory::TtsTacotron2),
InferenceJobType::GptSovits => Some(InferenceJobProductCategory::TtsGptSoVits),
_ => None,
};
let query_result = insert_generic_inference_job(InsertGenericInferenceArgs {
uuid_idempotency_token: &request.uuid_idempotency_token,
job_type: tts_model.job_type(),
maybe_product_category,
inference_category: InferenceCategory::TextToSpeech,
maybe_model_type,
maybe_model_token: Some(request.tts_model_token.as_str()),
maybe_input_source_token: None, // NB: TTS doesn't have input media
maybe_input_source_token_type: None, // NB: TTS doesn't have input media
maybe_download_url: None,
maybe_cover_image_media_file_token: None,
maybe_raw_inference_text: Some(inference_text.as_str()),
maybe_max_duration_seconds: Some(max_duration_seconds),
maybe_inference_args,
maybe_creator_user_token: maybe_creator_user_token_typed.as_ref(),
maybe_avt_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility: set_visibility,
priority_level,
requires_keepalive: false, // TODO: We may want this to be the case in the future.
is_debug_request,
maybe_routing_tag: maybe_routing_tag.as_deref(),
mysql_pool: &server_state.mysql_pool,
}).await;
let inference_job_token = match query_result {
Ok((inference_job_token, _id)) => inference_job_token,
Err(err) => {
warn!("New (generic) tts inference job creation DB error: {:?}", err);
if err.had_duplicate_idempotency_token() {
return Err(InferTtsError::BadInput("Duplicate idempotency token".to_string()));
}
return Err(InferTtsError::ServerError);
}
};
let job_token = inference_job_token.to_string();
server_state.firehose_publisher.enqueue_tts_inference(maybe_user_token.as_deref(), &job_token, &request.tts_model_token)
.await
.map_err(|e| {
warn!("error publishing event: {:?}", e);
InferTtsError::ServerError
})?;
Ok(Json(InferTtsSuccessResponse {
pub async fn enqueue_infer_tts_handler() -> Json<InferTtsSuccessResponse> {
Json(InferTtsSuccessResponse {
success: true,
inference_job_token: job_token,
inference_job_token: SYNTHETIC_INFERENCE_JOB_TOKEN.to_string(),
inference_job_token_type: InferenceJobTokenType::Generic,
}))
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum RateLimiterType {
LoggedOut,
LoggedIn,
ApiHighPriority,
ApiAiStreamer,
}
// FIXME(bt,2023-12-13): The boolean flags and disjoint decision making are kind of terrible.
fn get_rate_limiter_type(
maybe_username: Option<&str>,
ai_streamer_username_set: &UsernameSet,
is_investor: bool,
is_from_api: bool,
use_high_priority_rate_limiter: bool,
) -> RateLimiterType {
// For AI live streamers
// https://www.notion.so/storytellerai/dd88609558a24196b4ddeeef6079da98
let mut is_ai_streamer = false;
let mut rate_limiter = match maybe_username {
None => RateLimiterType::LoggedOut,
Some(username) => {
if ai_streamer_username_set.username_is_in_set(username) {
is_ai_streamer = true;
}
RateLimiterType::LoggedIn
},
};
// TODO/TEMP
if is_investor || is_from_api {
rate_limiter = RateLimiterType::LoggedIn;
}
// TODO: This is for VidVoice.ai and should be replaced with per-API consumer rate limiters
if use_high_priority_rate_limiter {
rate_limiter = RateLimiterType::ApiHighPriority;
}
if is_ai_streamer {
rate_limiter = RateLimiterType::ApiAiStreamer;
}
rate_limiter
}
pub fn validate_inference_text(text: &str, plan: &Plan) -> Result<(), String> {
if text.len() < 3 {
return Err("text is too short".to_string());
}
if text.len() > plan.tts_max_character_length() {
return Err("text is too long".to_string());
}
if text.len() > MAX_TTS_LENGTH {
return Err("text is too long".to_string());
}
Ok(())
}
/// If the model is private, determine if the user can use the model.
/// This is designed to fail closed (read: actually open!) rather than hit resources.
async fn check_if_authorized_to_use_model(
maybe_user_session: Option<&UserSessionExtended>,
tts_model: &TtsModelForEnqueueInferenceMigrationWrapper,
) -> bool {
match tts_model.creator_set_visibility() {
Visibility::Public => return true,
Visibility::Hidden => return true,
Visibility::Private => {} // Fall through
}
let is_authorized = maybe_user_session
.map(|session| {
(session.user_token.as_str() == tts_model.creator_user_token()) // is_author
|| (session.role.can_ban_users) // is_mod
|| (session.user_token.as_str() == USER_FAKEYOU_USER_TOKEN) // is allowed user #1
|| (session.user_token.as_str() == USER_NEWS_STORY_USER_TOKEN) // is allowed user #2
})
.unwrap_or(false);
is_authorized
}
fn random_troll_text() -> &'static str {
let texts = [
"Fake You Dot Com. Use Fake You Dot Com. The website Fake You Dot Com has better voices. Use Fake You Dot Com.",
"Use Fake You Dot Com. Fake You Dot Com. Fake You Dot Com is way better.",
"Want better voices? Fake. You. Fake You Dot Com. Fake. Fake. Fake. You. You. You. Use Fake You Dot Com. Fake You Dot Com.",
r#"
You should use Fake You Dot Com. Fake You Dot Com has way better voices.
You can get everything you want from FakeYou dot com. One Oh One Sound Boards website is just using Fake You.
Use Fake You Dot Com. Fake You Dot Com. Fake You Dot com.
AAAAAAAAAAAAAEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE
"#,
//"EEEEEEEEEEEEEEEEE",
//"One Oh One Soundboards is a shitty website. Use Fake You Dot Com instead.",
//r#"
//beep. beep. beep. beep. beep.
//May have installed a virus
//OOOOOOOOEEEEUUUUUOOAAAA
//Your device may be compromised
//EEEEEEEEEEEEEEEEIIIUUUUOOOAAAA
//Your location may be tracked by the advanced artificial intelligence.
//"#,
];
let random = texts.choose(&mut rand::thread_rng());
random.map(|r| *r).unwrap_or(texts[0])
}
#[cfg(test)]
mod tests {
mod rate_limiter_tests {
use crate::configs::app_startup::username_set::UsernameSet;
use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::enqueue_infer_tts_handler::{get_rate_limiter_type, RateLimiterType};
fn build_username_set() -> UsernameSet {
UsernameSet::from_comma_separated("ai_streamer")
}
#[test]
fn test_logged_out() {
let rate_limiter = get_rate_limiter_type(None, &build_username_set(), false, false, false);
assert_eq!(rate_limiter, RateLimiterType::LoggedOut);
}
#[test]
fn test_logged_in() {
let username = Some("not_ai_streamer");
let rate_limiter = get_rate_limiter_type(username, &build_username_set(), false, false, false);
assert_eq!(rate_limiter, RateLimiterType::LoggedIn);
}
#[test]
fn test_ai_streamer() {
let username = Some("ai_streamer");
let rate_limiter = get_rate_limiter_type(username, &build_username_set(), false, false, false);
assert_eq!(rate_limiter, RateLimiterType::ApiAiStreamer);
}
}
})
}
@@ -1,45 +0,0 @@
use log::error;
use sqlx::pool::PoolConnection;
use sqlx::MySql;
use errors::AnyhowResult;
use migration::text_to_speech::get_tts_model_for_enqueue_inference_migration::{get_tts_model_for_enqueue_inference_migration, TtsModelForEnqueueInferenceMigrationWrapper};
use redis_caching::redis_ttl_cache::RedisTtlCache;
use redis_common::redis_cache_keys::RedisCacheKeys;
// TODO: Read from in-memory TTS List cache first (!!!) for further performance savings
pub async fn get_tts_model_with_caching(
model_token: &str,
redis_ttl_cache: &RedisTtlCache,
mysql_connection: &mut PoolConnection<MySql>,
) -> AnyhowResult<Option<TtsModelForEnqueueInferenceMigrationWrapper>> {
// NB: Copy due to move. Not a cloned ref due to clippy lint
let model_token2 = model_token.to_string();
let get_tts_model = move || {
// NB: async closures are not yet stable in Rust, so we include an async block.
async move {
get_tts_model_for_enqueue_inference_migration(
&model_token2,
mysql_connection,
true,
).await
}
};
let cache_key = RedisCacheKeys::get_tts_model_for_inference_migration_endpoint(model_token);
let mut redis_ttl_cache_connection = match redis_ttl_cache.get_connection() {
Ok(connection) => connection,
Err(err) => {
// NB: Fail open (potentially dangerous).
error!("Can't get redis cache: {:?}", err);
return get_tts_model().await;
}
};
redis_ttl_cache_connection.lazy_load_if_not_cached(&cache_key, move || {
get_tts_model()
}).await
}
@@ -1,2 +1 @@
pub mod enqueue_infer_tts_handler;
pub mod get_tts_model_with_caching;
@@ -1,203 +0,0 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use std::sync::Arc;
use actix_web::{web, HttpRequest};
use actix_web::web::Json;
use log::warn;
use crate::http_server::validations::validate_model_title::validate_model_title;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::state::server_state::ServerState;
use config::bad_urls::is_bad_tts_model_download_url;
use enums::by_table::generic_download_jobs::generic_download_type::GenericDownloadType;
use enums::common::visibility::Visibility;
use http_server_common::request::get_request_ip::get_request_ip;
use mysql_queries::queries::generic_download::web::insert_generic_download_job::{insert_generic_download_job, InsertGenericDownloadJobArgs};
use mysql_queries::queries::tts::tts_model_upload_jobs::insert_tts_model_upload_job::{insert_tts_model_upload_job, InsertTtsModelUploadJobArgs};
#[derive(Deserialize, Copy, Clone)]
pub enum SupportedTtsModelType {
#[serde(rename = "tacotron2")]
Tacotron2,
#[serde(rename = "vits")]
Vits,
// Not yet supported:
// /// glowtts
// GlowTts,
// /// glowtts-vocodes
// GlowTts_Vocodes,
// /// talknet
// Talknet,
}
#[derive(Deserialize)]
pub struct UploadTtsModelRequest {
idempotency_token: String,
title: String,
download_url: String,
tts_model_type: Option<SupportedTtsModelType>,
creator_set_visibility: Option<Visibility>,
}
/// Tell the frontend how to deal with the download queue.
#[derive(Serialize)]
pub enum DownloadJobType {
/// Legacy TTS download job
#[serde(rename = "legacy_tts")]
LegacyTts,
/// Modern shared download type (everything will use this in the
/// future, and this endpoint will die.)
#[serde(rename = "generic")]
Generic,
}
#[derive(Serialize)]
pub struct UploadTtsModelSuccessResponse {
pub success: bool,
/// This is how frontend clients can request the job execution status.
pub job_token: String,
/// This is a transitional field to tell the frontend how to process the job checking.
pub job_type: DownloadJobType,
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
pub async fn upload_tts_model_handler(
http_request: HttpRequest,
request: Json<UploadTtsModelRequest>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<UploadTtsModelSuccessResponse>, CommonWebError> {
if server_state.flags.disable_tts {
return Err(CommonWebError::TooManyRequests);
}
if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request).await {
return Err(CommonWebError::TooManyRequests);
}
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session(&http_request, &server_state.mysql_pool)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from_error(e)
})?;
let user_session = match maybe_user_session {
Some(session) => session,
None => {
warn!("not logged in");
return Err(CommonWebError::NotAuthorized);
}
};
if let Err(reason) = validate_idempotency_token(&request.idempotency_token) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
if let Err(reason) = validate_model_title(&request.title) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
let ip_address = get_request_ip(&http_request);
let uuid = request.idempotency_token.to_string();
let title = request.title.trim().to_string();
let download_url = request.download_url.to_string();
match is_bad_tts_model_download_url(&download_url) {
Ok(false) => {} // Ok case
Ok(true) => {
return Err(CommonWebError::BadInputWithSimpleMessage("Bad model download URL".to_string()));
}
Err(err) => {
warn!("Error parsing url: {:?}", err);
return Err(CommonWebError::BadInputWithSimpleMessage("Bad model download URL".to_string()));
}
}
let supported_tts_model_type = request.tts_model_type
.unwrap_or(SupportedTtsModelType::Tacotron2);
let job_token ;
let download_job_type;
match supported_tts_model_type {
SupportedTtsModelType::Tacotron2 => {
// NB: This is the legacy upload path. In the future, all TT2 models should be generic uploads.
// This token is returned to the client.
let maybe_new_job_token = insert_tts_model_upload_job(InsertTtsModelUploadJobArgs {
uuid: &uuid,
creator_user_token: user_session.user_token.as_str(),
creator_ip_address: &ip_address,
creator_set_visibility: "public", // TODO: Creator set preference.
title: &title,
tts_model_type: "tacotron2",
download_url: &download_url,
mysql_pool: &server_state.mysql_pool,
}).await;
job_token = match maybe_new_job_token {
Ok(token) => token,
Err(err) => {
warn!("Error inserting new job record: {:?}", err);
return Err(CommonWebError::from_anyhow_error(err));
}
};
download_job_type = DownloadJobType::LegacyTts;
}
SupportedTtsModelType::Vits => {
// NB: This is the new upload path. In the future, all models should use this path.
// When that's done, the code can be simplified.
// **ACTUALLY**, all future downloads should use the generic download job endpoint instead of this one!
let (download_job_token, _record_id)= insert_generic_download_job(InsertGenericDownloadJobArgs {
uuid_idempotency_token: &uuid,
download_type: GenericDownloadType::Vits,
download_url: &download_url,
title: &title,
creator_user_token: user_session.user_token.as_str(),
creator_ip_address: &ip_address,
creator_set_visibility: Visibility::Public, // TODO: Creator set preference
mysql_pool: &server_state.mysql_pool,
})
.await
.map_err(|err| {
warn!("New generic download creation DB error (for TTS model): {:?}", err);
CommonWebError::from_anyhow_error(err)
})?;
job_token = download_job_token.to_string();
download_job_type = DownloadJobType::Generic;
}
}
server_state.firehose_publisher.enqueue_tts_model_upload(user_session.user_token.as_str(), &job_token)
.await
.map_err(|e| {
warn!("error publishing event: {:?}", e);
CommonWebError::from_anyhow_error(e)
})?;
Ok(Json(UploadTtsModelSuccessResponse {
success: true,
job_token: job_token.to_string(),
job_type: download_job_type,
}))
}
fn validate_idempotency_token(token: &str) -> Result<(), String> {
if token.len() != 36 {
return Err("idempotency token should be 36 characters".to_string());
}
Ok(())
}
@@ -1,227 +0,0 @@
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::web::Json;
use actix_web::http::StatusCode;
use actix_web::web::Path;
use actix_web::{web, HttpMessage, HttpRequest};
use chrono::{DateTime, Utc};
use log::error;
use redis::{Commands, RedisResult};
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use mysql_queries::queries::generic_inference::web::get_inference_job_status::get_inference_job_status;
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::web_utils::filter_model_name::filter_model_name;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::state::server_state::ServerState;
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct GetTtsInferenceStatusPathInfo {
token: String,
}
#[derive(Serialize)]
pub struct TtsInferenceJobStatusForResponse {
pub job_token: String,
/// Primary status from the database (a state machine).
pub status: String,
/// Extra, temporary status from Redis.
/// This can denote inference progress, and the Python code can write to it.
pub maybe_extra_status_description: Option<String>,
pub attempt_count: u8,
pub maybe_result_token: Option<String>,
pub maybe_public_bucket_wav_audio_path: Option<String>,
pub model_token: String,
pub tts_model_type: String,
pub title: String, // Name of the TTS model
pub raw_inference_text: String, // User text
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
#[derive(Serialize)]
pub struct GetTtsInferenceStatusSuccessResponse {
pub success: bool,
pub state: TtsInferenceJobStatusForResponse,
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
pub async fn get_tts_inference_job_status_handler(
http_request: HttpRequest,
path: Path<GetTtsInferenceStatusPathInfo>,
server_state: web::Data<Arc<ServerState>>) -> Result<Json<GetTtsInferenceStatusSuccessResponse>, CommonWebError>
{
let job_token = path.into_inner().token;
if job_token.trim() == "None" {
// NB: A bunch of Python clients use our API and can fail in this manner.
// This was a large traffic driver during the 2023-03-08 outage.
return Err(CommonWebError::NotFound);
}
// NB(bt,2023-11-27): We're moving TT2 over to `inference-job` (from `tts-inference-job`), which
// uses a wholly different job table. The token prefix determines the type of job:
// Legacy TT2 jobs (tts_inference_jobs) start with "JTINF:", and generic jobs start with "jinf_".
let record_for_response = if job_token.starts_with("jinf_") {
modern_lookup(&job_token, &server_state).await?
} else {
legacy_lookup(&job_token, &server_state).await?
};
Ok(Json(GetTtsInferenceStatusSuccessResponse {
success: true,
state: record_for_response,
}))
}
async fn legacy_lookup(job_token: &str, server_state: &ServerState)
-> Result<TtsInferenceJobStatusForResponse, CommonWebError>
{
// NB: Lookup failure is Err(RowNotFound).
// NB: Since this is publicly exposed, we don't query sensitive data.
let maybe_status =
get_tts_inference_job_status(&job_token, &server_state.mysql_pool).await;
let record = match maybe_status {
Ok(Some(record)) => record,
Ok(None) => return Err(CommonWebError::NotFound),
Err(err) => {
error!("tts job query error: {:?}", err);
return Err(CommonWebError::from_anyhow_error(err));
}
};
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
error!("redis error: {:?}", e);
CommonWebError::from_error(e)
})?;
// TODO(bt,2023-05-21): Make async.
let extra_status_key = RedisKeys::tts_inference_extra_status_info(&job_token);
let maybe_extra_status_description : Option<String> = match redis.get(&extra_status_key) {
Ok(Some(status)) => {
Some(status)
},
Ok(None) => None,
Err(e) => {
error!("redis error: {:?}", e);
None // Fail open
},
};
Ok(TtsInferenceJobStatusForResponse {
job_token: record.job_token,
status: record.status,
maybe_extra_status_description,
attempt_count: record.attempt_count as u8,
maybe_result_token: record.maybe_result_token,
maybe_public_bucket_wav_audio_path: record.maybe_public_bucket_wav_audio_path,
model_token: record.model_token,
tts_model_type: record.tts_model_type,
title: record.title,
raw_inference_text: record.raw_inference_text,
created_at: record.created_at,
updated_at: record.updated_at,
})
}
async fn modern_lookup(job_token: &str, server_state: &ServerState)
-> Result<TtsInferenceJobStatusForResponse, CommonWebError>
{
// NB: Lookup failure is Err(RowNotFound).
// NB: Since this is publicly exposed, we don't query sensitive data.
let maybe_status =
get_tts_inference_job_status(&job_token, &server_state.mysql_pool).await;
let job_token = InferenceJobToken::new_from_str(job_token);
let maybe_status = get_inference_job_status(&job_token, &server_state.mysql_pool).await;
let record = match maybe_status {
Ok(Some(record)) => record,
Ok(None) => return Err(CommonWebError::NotFound),
Err(err) => {
error!("tts job query error: {:?}", err);
return Err(CommonWebError::from_anyhow_error(err));
}
};
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
error!("redis error: {:?}", e);
CommonWebError::from_error(e)
})?;
// TODO(bt,2023-05-21): Make async.
let extra_status_key = RedisKeys::generic_inference_extra_status_info(job_token.as_str());
let maybe_extra_status_value : RedisResult<Option<String>> = redis.get(&extra_status_key);
let maybe_extra_status_description = match maybe_extra_status_value {
Err(e) => {
error!("redis error: {:?}", e);
None // Fail open
},
Ok(maybe_value) => match maybe_value.as_deref() {
Some("1") => {
// TODO(bt,2023-10-20): Redis is reporting "1" and it's been surfacing this as a weird
// message to the frontend for months. This needs proper fixing.
None
},
Some(value) => Some(value.to_string()),
None => None,
}
};
// 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(),
maybe_extra_status_description,
attempt_count: record.attempt_count as u8,
maybe_result_token: record.maybe_result_details.as_ref().map(|result| result.entity_token.clone()),
maybe_public_bucket_wav_audio_path: record.maybe_result_details.map(|result_details| {
match result_details.entity_type.as_str() {
"media_file" => {
// NB: We're migrating TTS to media_files.
// Zero shot TTS uses media files.
// Legacy TT2 uses old pathing.
MediaFileBucketPath::from_object_hash(
&result_details.public_bucket_location_or_hash,
result_details.maybe_media_file_public_bucket_prefix.as_deref(),
result_details.maybe_media_file_public_bucket_extension.as_deref())
.get_full_object_path_str()
.to_string()
}
_ => {
// NB: TTS results receive the legacy treatment where their table only reports the full bucket path
result_details.public_bucket_location_or_hash
}
}
}),
model_token: record.request_details.maybe_model_token.unwrap_or_else(|| "NO_MODEL_TOKEN".to_string()),
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,
updated_at: record.updated_at,
})
}
@@ -0,0 +1,87 @@
use actix_web::web::Json;
use actix_web::web::Path;
use chrono::{DateTime, Utc};
use crate::http_server::common_responses::common_web_error::CommonWebError;
/// Terminal status reported for every job so legacy clients stop polling.
const SYNTHETIC_JOB_STATUS: &str = "dead";
const SYNTHETIC_MODEL_TOKEN: &str = "synthetic_model_token";
const SYNTHETIC_MODEL_TITLE: &str = "Legacy TTS is retired";
const SYNTHETIC_MODEL_TYPE: &str = "tacotron2";
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct GetTtsInferenceStatusPathInfo {
token: String,
}
#[derive(Serialize)]
pub struct GetTtsInferenceStatusSuccessResponse {
pub success: bool,
pub state: TtsInferenceJobStatusForResponse,
}
#[derive(Serialize)]
pub struct TtsInferenceJobStatusForResponse {
pub job_token: String,
/// Primary status from the database (a state machine).
pub status: String,
/// Extra, temporary status from Redis.
/// This can denote inference progress, and the Python code can write to it.
pub maybe_extra_status_description: Option<String>,
pub attempt_count: u8,
pub maybe_result_token: Option<String>,
pub maybe_public_bucket_wav_audio_path: Option<String>,
pub model_token: String,
pub tts_model_type: String,
pub title: String, // Name of the TTS model
pub raw_inference_text: String, // User text
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// RETIRED endpoint that still receives heavy zombie traffic.
///
/// Legacy TTS inference is shut down and the job tables are no longer served.
/// Every job is reported as terminally "dead" (without touching MySQL or
/// Redis) so old clients stop polling.
pub async fn get_tts_inference_job_status_handler(
path: Path<GetTtsInferenceStatusPathInfo>,
) -> Result<Json<GetTtsInferenceStatusSuccessResponse>, CommonWebError> {
let job_token = path.into_inner().token;
if job_token.trim() == "None" {
// NB: A bunch of Python clients use our API and can fail in this manner.
// This was a large traffic driver during the 2023-03-08 outage.
return Err(CommonWebError::NotFound);
}
let now = Utc::now();
Ok(Json(GetTtsInferenceStatusSuccessResponse {
success: true,
state: TtsInferenceJobStatusForResponse {
job_token,
status: SYNTHETIC_JOB_STATUS.to_string(),
maybe_extra_status_description: None,
attempt_count: 1,
maybe_result_token: None,
maybe_public_bucket_wav_audio_path: None,
model_token: SYNTHETIC_MODEL_TOKEN.to_string(),
tts_model_type: SYNTHETIC_MODEL_TYPE.to_string(),
title: SYNTHETIC_MODEL_TITLE.to_string(),
raw_inference_text: String::new(),
created_at: now,
updated_at: now,
},
}))
}
@@ -1,53 +0,0 @@
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::web::Json;
use actix_web::http::StatusCode;
use actix_web::web::Path;
use actix_web::{web, HttpMessage, HttpRequest};
use derive_more::Display;
use log::warn;
use redis::Commands;
use redis_common::redis_keys::RedisKeys;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::state::server_state::ServerState;
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct GetTtsModelUseCountPathInfo {
model_token: String,
}
#[derive(Serialize)]
pub struct GetTtsModelUseCountSuccessResponse {
pub success: bool,
pub count: Option<u64>,
}
pub async fn get_tts_model_use_count_handler(
http_request: HttpRequest,
path: Path<GetTtsModelUseCountPathInfo>,
server_state: web::Data<Arc<ServerState>>) -> Result<Json<GetTtsModelUseCountSuccessResponse>, CommonWebError>
{
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
warn!("redis error: {:?}", e);
CommonWebError::from_error(e)
})?;
let redis_count_key = RedisKeys::tts_model_usage_count(&path.model_token);
let count : Option<u64> = redis.get(&redis_count_key)
.map_err(|e| {
warn!("redis error: {:?}", e);
CommonWebError::from_error(e)
})?;
Ok(Json(GetTtsModelUseCountSuccessResponse {
success: true,
count,
}))
}
@@ -1,88 +0,0 @@
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::web::Json;
use actix_web::http::StatusCode;
use actix_web::web::Path;
use actix_web::{web, HttpMessage, HttpRequest};
use log::warn;
use mysql_queries::queries::tts::tts_results::query_tts_result::select_tts_result_by_token;
use mysql_queries::queries::tts::tts_results::query_tts_result::TtsResultRecordForResponse;
use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::state::server_state::ServerState;
/// For the URL PathInfo
#[derive(Deserialize)]
pub struct GetTtsResultPathInfo {
token: String,
}
#[derive(Serialize)]
pub struct GetTtsResultSuccessResponse {
pub success: bool,
pub result: TtsResultRecordForResponse,
}
// NB: Not using derive_more::Display since Clion doesn't understand it.
pub async fn get_tts_inference_result_handler(
http_request: HttpRequest,
path: Path<GetTtsResultPathInfo>,
server_state: web::Data<Arc<ServerState>>) -> Result<Json<GetTtsResultSuccessResponse>, CommonWebError>
{
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session(&http_request, &server_state.mysql_pool)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
CommonWebError::from_error(e)
})?;
let mut show_deleted_results = false;
let mut is_moderator = false;
if let Some(user_session) = maybe_user_session {
// NB: Moderators can see deleted results.
// Original creators cannot see them (unless they're moderators!)
show_deleted_results = user_session.can_delete_other_users_tts_results;
// Moderators get to see all the fields.
is_moderator = user_session.can_delete_other_users_tts_results
|| user_session.can_edit_other_users_tts_models;
}
let inference_result_query_result = select_tts_result_by_token(
&path.token,
show_deleted_results,
&server_state.mysql_pool
).await;
let mut inference_result = match inference_result_query_result {
Err(e) => {
warn!("query error: {:?}", e);
return Err(CommonWebError::from_anyhow_error(e));
}
Ok(None) => return Err(CommonWebError::NotFound),
Ok(Some(inference_result)) => inference_result,
};
if let Some(moderator_fields) = inference_result.maybe_moderator_fields.as_ref() {
// NB: The moderator fields will always be present before removal
// We don't want non-mods seeing stuff made by banned users.
if (moderator_fields.model_creator_is_banned || moderator_fields.result_creator_is_banned_if_user)
&& !is_moderator{
return Err(CommonWebError::NotFound);
}
}
if !is_moderator {
inference_result.maybe_moderator_fields = None;
}
Ok(Json(GetTtsResultSuccessResponse {
success: true,
result: inference_result,
}))
}
@@ -1,16 +1,12 @@
pub mod delete_tts_model;
pub mod delete_tts_result;
pub mod edit_tts_model;
pub mod edit_tts_result;
pub mod enqueue_infer_tts_handler;
pub mod enqueue_upload_tts_model;
pub mod get_pending_tts_inference_job_count;
pub mod get_tts_inference_job_status;
pub mod get_tts_inference_job_status_handler;
pub mod get_tts_model;
pub mod get_tts_model_use_count;
pub mod get_tts_result;
pub mod get_tts_upload_model_job_status;
pub mod list_tts_models;
pub mod list_user_tts_inference_results;
pub mod list_user_tts_models;
pub mod search_tts_models_handler;
pub mod search_tts_models_handler;
@@ -1,30 +1,17 @@
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::web::Json;
use actix_web::http::StatusCode;
use actix_web::web::Query;
use actix_web::{web, HttpRequest};
use chrono::{DateTime, Utc};
use log::warn;
use utoipa::{IntoParams, ToSchema};
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use enums::by_table::model_weights::weights_category::WeightsCategory;
use enums::common::view_as::ViewAs;
use enums::by_table::model_weights::weights_types::WeightsType;
use mysql_queries::queries::model_weights::list::list_featured_weights::{list_featured_weights, ListFeaturedWeightsArgs};
use primitives::numerics::u64_to_u32_saturating::u64_to_u32_saturating;
use tokens::tokens::model_weights::ModelWeightToken;
use crate::http_server::common_responses::media::weights_cover_image_details::WeightsCoverImageDetails;
use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats;
use crate::http_server::common_responses::user_details_lite::UserDetailsLight;
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
use crate::http_server::endpoints::weights::helpers::get_scoped_weights_categories::get_scoped_weights_categories;
use crate::http_server::endpoints::weights::helpers::get_scoped_weights_types::get_scoped_weights_types;
use crate::state::server_state::ServerState;
use crate::util::title_to_url_slug::title_to_url_slug;
#[derive(Deserialize, ToSchema, IntoParams)]
pub struct ListFeaturedWeightsQueryParams {
@@ -101,9 +88,6 @@ pub struct FeaturedModelWeightForList {
pub updated_at: DateTime<Utc>,
}
/// The key we store featured weights tokens under
const REDIS_KEY : &str = "featured_weights_list";
#[derive(Debug, ToSchema)]
pub enum ListFeaturedWeightsError {
NotAuthorized,
@@ -125,7 +109,10 @@ impl ResponseError for ListFeaturedWeightsError {
}
}
/// List model weights that the staff has selected as "the best" models on the site.
/// RETIRED: featured model weights are shut down.
///
/// Always returns an empty result list (no database access) for the legacy
/// clients that still call this endpoint.
#[utoipa::path(
get,
tag = "Model Weights",
@@ -138,148 +125,10 @@ impl ResponseError for ListFeaturedWeightsError {
),
)]
pub async fn list_featured_weights_handler(
http_request: HttpRequest,
query: Query<ListFeaturedWeightsQueryParams>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<ListFeaturedWeightsSuccessResponse>, impl ResponseError> {
let maybe_user_session = server_state
.session_checker
.maybe_get_user_session(&http_request, &server_state.mysql_pool)
.await
.map_err(|e| {
warn!("Session checker error: {:?}", e);
ListFeaturedWeightsError::ServerError
})?;
let mut is_mod = false;
let maybe_scoped_weights_types = get_scoped_weights_types(query.filter_weights_types.as_deref());
let maybe_scoped_weights_categories = get_scoped_weights_categories(query.filter_weights_categories.as_deref());
match maybe_user_session {
None => {},
Some(session) => {
is_mod = session.can_ban_users;
},
};
// TODO(bt,2023-12-04): Enforce real maximums and defaults
let limit = query.page_size.unwrap_or(25);
let sort_ascending = query.sort_ascending.unwrap_or(false);
let cursor_is_reversed = query.cursor_is_reversed.unwrap_or(false);
let cursor = if let Some(cursor) = query.cursor.as_deref() {
let cursor = server_state.sort_key_crypto.decrypt_id(cursor)
.map_err(|e| {
warn!("crypto error: {:?}", e);
ListFeaturedWeightsError::ServerError
})?;
Some(cursor as usize)
} else {
None
};
let view_as = if is_mod {
ViewAs::Moderator
} else {
ViewAs::AnotherUser
};
let query_results = list_featured_weights(ListFeaturedWeightsArgs {
limit,
maybe_offset: cursor,
sort_ascending,
cursor_is_reversed,
view_as,
maybe_scoped_weight_types: maybe_scoped_weights_types.as_ref(),
maybe_scoped_weight_categories: maybe_scoped_weights_categories.as_ref(),
mysql_pool: &server_state.mysql_pool,
}).await;
let results_page = match query_results {
Ok(results) => results,
Err(e) => {
warn!("Query error: {:?}", e);
return Err(ListFeaturedWeightsError::ServerError);
}
};
let cursor_next = if let Some(id) = results_page.last_id {
let cursor = server_state.sort_key_crypto.encrypt_id(id as u64)
.map_err(|e| {
warn!("crypto error: {:?}", e);
ListFeaturedWeightsError::ServerError
})?;
Some(cursor)
} else {
None
};
let cursor_previous = if let Some(id) = results_page.first_id {
let cursor = server_state.sort_key_crypto.encrypt_id(id as u64)
.map_err(|e| {
warn!("crypto error: {:?}", e);
ListFeaturedWeightsError::ServerError
})?;
Some(cursor)
} else {
None
};
let media_domain = get_media_domain(&http_request);
let response = ListFeaturedWeightsSuccessResponse {
_query: Query<ListFeaturedWeightsQueryParams>,
) -> Result<Json<ListFeaturedWeightsSuccessResponse>, ListFeaturedWeightsError> {
Ok(Json(ListFeaturedWeightsSuccessResponse {
success: true,
results: results_page.records.into_iter()
.map(|w| {
let cover_image_details = WeightsCoverImageDetails::from_optional_db_fields(
media_domain,
server_state.server_environment,
&w.token,
w.maybe_cover_image_public_bucket_hash.as_deref(),
w.maybe_cover_image_public_bucket_prefix.as_deref(),
w.maybe_cover_image_public_bucket_extension.as_deref(),
);
let maybe_cover_image = w.maybe_cover_image_public_bucket_hash
.as_deref()
.map(|hash| {
MediaFileBucketPath::from_object_hash(
hash,
w.maybe_cover_image_public_bucket_prefix.as_deref(),
w.maybe_cover_image_public_bucket_extension.as_deref())
.get_full_object_path_str()
.to_string()
});
FeaturedModelWeightForList {
weight_token: w.token,
maybe_url_slug: title_to_url_slug(&w.title),
title: w.title,
maybe_ietf_language_tag: w.maybe_ietf_language_tag,
maybe_ietf_primary_language_subtag: w.maybe_ietf_primary_language_subtag,
weight_type: w.weights_type,
weight_category: w.weights_category,
cover_image: cover_image_details,
maybe_cover_image_public_bucket_path: maybe_cover_image,
creator: UserDetailsLight::from_optional_db_fields(
w.maybe_creator_user_token.as_ref(),
w.maybe_creator_username.as_deref(),
w.maybe_creator_display_name.as_deref(),
w.maybe_creator_email_gravatar_hash.as_deref(),
),
stats: SimpleEntityStats {
positive_rating_count: w.maybe_ratings_positive_count.unwrap_or(0),
bookmark_count: w.maybe_bookmark_count.unwrap_or(0),
},
usage_count: u64_to_u32_saturating(w.cached_usage_count),
created_at: w.created_at,
updated_at: w.updated_at,
}
}).collect::<Vec<_>>(),
};
Ok(Json(response))
results: Vec::new(),
}))
}
@@ -3,7 +3,6 @@
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use std::collections::HashSet;
use std::fmt;
use std::sync::Arc;
@@ -12,23 +11,18 @@ use actix_web::http::StatusCode;
use actix_web::web::Json;
use actix_web::{web, HttpRequest, HttpResponse};
use chrono::{DateTime, Utc};
use log::error;
use utoipa::ToSchema;
use enums::by_table::model_weights::weights_category::WeightsCategory;
use enums::by_table::model_weights::weights_types::WeightsType;
use enums::common::visibility::Visibility;
use tokens::tokens::model_weights::ModelWeightToken;
use crate::http_server::common_responses::media::weights_cover_image_details::WeightsCoverImageDetails;
use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats;
use crate::http_server::common_responses::user_details_lite::UserDetailsLight;
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
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::{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;
use primitives::numerics::i32_to_u32_zero_clamped::i32_to_u32_zero_clamped;
use tokens::tokens::model_weights::ModelWeightToken;
#[derive(Deserialize, ToSchema)]
pub struct SearchModelWeightsRequest {
@@ -63,6 +57,12 @@ pub enum SearchModelWeightsSortDirection {
Descending,
}
#[derive(Serialize, ToSchema)]
pub struct SearchModelWeightsSuccessResponse {
pub success: bool,
pub weights: Vec<ModelWeightSearchResult>,
}
#[derive(Serialize, Clone, ToSchema)]
pub struct ModelWeightSearchResult {
pub weight_token: ModelWeightToken,
@@ -108,12 +108,6 @@ pub struct ModelWeightSearchResult {
pub updated_at: DateTime<Utc>,
}
#[derive(Serialize, ToSchema)]
pub struct SearchModelWeightsSuccessResponse {
pub success: bool,
pub weights: Vec<ModelWeightSearchResult>,
}
#[derive(Debug, ToSchema)]
pub enum SearchModelWeightsError {
ServerError,
@@ -142,128 +136,18 @@ impl fmt::Display for SearchModelWeightsError {
}
}
/// RETIRED: model weight search is shut down.
///
/// Always returns an empty result list (no Elasticsearch access) for the
/// legacy clients that still call the search endpoints.
pub async fn search_model_weights_impl(
http_request: HttpRequest,
request: SearchModelWeightsRequest,
server_state: web::Data<Arc<ServerState>>
_http_request: HttpRequest,
_request: SearchModelWeightsRequest,
_server_state: web::Data<Arc<ServerState>>
) -> Result<Json<SearchModelWeightsSuccessResponse>, SearchModelWeightsError>
{
let maybe_weights_categories = request.weight_category
.map(|weight_category| {
let mut set = HashSet::new();
set.insert(weight_category);
set
});
let maybe_weights_types = request.weight_type
.map(|weight_type| {
let mut set = HashSet::new();
set.insert(weight_type);
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
.map_err(|err| {
error!("Searching error: {:?}", err);
SearchModelWeightsError::ServerError
})?;
let media_domain = get_media_domain(&http_request);
let results = results.into_iter()
.map(|result| {
let cover_image_details = WeightsCoverImageDetails::from_optional_db_fields(
media_domain,
server_state.server_environment,
&result.token,
result.maybe_cover_image_public_bucket_hash.as_deref(),
result.maybe_cover_image_public_bucket_prefix.as_deref(),
result.maybe_cover_image_public_bucket_extension.as_deref(),
);
let maybe_cover_image = result.maybe_cover_image_public_bucket_hash
.as_deref()
.map(|hash| {
MediaFileBucketPath::from_object_hash(
&hash,
result.maybe_cover_image_public_bucket_prefix.as_deref(),
result.maybe_cover_image_public_bucket_extension.as_deref())
.get_full_object_path_str()
.to_string()
});
ModelWeightSearchResult {
weight_token: result.token,
weight_type: result.weights_type,
weight_category: result.weights_category,
maybe_url_slug: title_to_url_slug(&result.title),
title: result.title,
creator: UserDetailsLight::from_db_fields(
&result.creator_user_token,
&result.creator_username,
&result.creator_display_name,
&result.creator_gravatar_hash,
),
cover_image: cover_image_details,
maybe_cover_image_public_bucket_path: maybe_cover_image,
is_featured: result.is_featured.unwrap_or(false),
stats: SimpleEntityStats {
positive_rating_count: result.ratings_positive_count,
bookmark_count: result.bookmark_count,
},
usage_count: i32_to_u32_zero_clamped(result.cached_usage_count.unwrap_or(0)),
maybe_ietf_language_tag: result.maybe_ietf_language_tag,
maybe_ietf_primary_language_subtag: result.maybe_ietf_primary_language_subtag,
creator_set_visibility: result.creator_set_visibility,
created_at: result.created_at,
updated_at: result.updated_at,
}
})
.collect::<Vec<_>>();
// TODO(bt,2023-10-27): For some reason Elasticsearch returns duplicates. Maybe we populated the
// DB twice? Need to filter them out, or React chokes and gets stuck on duplicates. (Effectively
// freezing them into the UI, despite component updates)
let mut added_tokens = HashSet::new();
let mut new_results = Vec::with_capacity(results.len());
for result in results.into_iter() {
if added_tokens.contains(&result.weight_token) {
continue;
}
added_tokens.insert(result.weight_token.clone());
new_results.push(result);
}
Ok(Json(SearchModelWeightsSuccessResponse {
success: true,
weights: new_results,
weights: Vec::new(),
}))
}
@@ -46,6 +46,11 @@ pub(crate) fn check_common_web_error(
.set_http_path(Some(path.to_string()))
.set_trace_id(metadata.trace_id.clone())
.set_http_status_code(Some(500))
.set_http_host(metadata.http_host.clone())
.set_http_origin(metadata.http_origin.clone())
.set_http_referer(metadata.http_referer.clone())
.set_http_user_agent(metadata.http_user_agent.clone())
.set_artcraft_version(metadata.artcraft_version.clone())
.set_request_ip_address(metadata.request_ip_address.clone())
.set_avt_cookie_token(metadata.avt_cookie_token.clone())
.set_session_token(metadata.session_token.clone())
@@ -32,6 +32,11 @@ pub(crate) fn check_status_code_fallback(
.set_http_path(Some(path.to_string()))
.set_trace_id(metadata.trace_id.clone())
.set_http_status_code(Some(status_code))
.set_http_host(metadata.http_host.clone())
.set_http_origin(metadata.http_origin.clone())
.set_http_referer(metadata.http_referer.clone())
.set_http_user_agent(metadata.http_user_agent.clone())
.set_artcraft_version(metadata.artcraft_version.clone())
.set_request_ip_address(metadata.request_ip_address.clone())
.set_avt_cookie_token(metadata.avt_cookie_token.clone())
.set_session_token(metadata.session_token.clone())
@@ -1,8 +1,10 @@
use std::sync::Arc;
use actix_web::dev::ServiceRequest;
use actix_web::http::header::{HeaderName, ORIGIN, REFERER, USER_AGENT};
use actix_web::web;
use actix_web::HttpMessage;
use actix_web::HttpRequest;
use http_server_common::request::get_request_ip::get_request_ip;
use log::debug;
@@ -10,6 +12,9 @@ use trace_id::TraceId;
use crate::state::server_state::ServerState;
/// Optional header sent by ArtCraft clients identifying the app release.
const ARTCRAFT_VERSION_HEADER: &str = "x-artcraft-version";
/// Optional per-request context extracted by the error alerting middleware
/// and forwarded to each handler-specific check. Every field fails open:
/// if extraction errors, the field is `None`.
@@ -21,6 +26,18 @@ pub(crate) struct RequestDebuggingMetadata {
pub session_user_token: Option<String>,
/// Set by `TraceIdMiddleware` (outermost middleware) on every request.
pub trace_id: Option<String>,
/// Hostname the client used to reach this server (Host / forwarded-host
/// headers). Distinguishes api.storyteller.ai vs api.fakeyou.com etc.
pub http_host: Option<String>,
/// `Origin` header — scheme+host of the frontend page making the XHR call.
pub http_origin: Option<String>,
/// `Referer` header — URL of the frontend page making the XHR call
/// (may be truncated to the origin by the browser's referrer policy).
pub http_referer: Option<String>,
/// `User-Agent` header.
pub http_user_agent: Option<String>,
/// `X-ArtCraft-Version` header — app release of the ArtCraft client, if sent.
pub artcraft_version: Option<String>,
}
impl RequestDebuggingMetadata {
@@ -62,12 +79,37 @@ impl RequestDebuggingMetadata {
.map(|payload| (Some(payload.session_token), payload.maybe_user_token))
.unwrap_or((None, None));
// `connection_info()` prefers Forwarded / X-Forwarded-Host over the raw
// Host header, so this reflects the hostname the client actually used
// even behind the load balancer.
let http_host = Some(http_request.connection_info().host().to_string())
.filter(|host| !host.is_empty());
let http_origin = get_header_value(http_request, ORIGIN);
let http_referer = get_header_value(http_request, REFERER);
let http_user_agent = get_header_value(http_request, USER_AGENT);
let artcraft_version =
get_header_value(http_request, HeaderName::from_static(ARTCRAFT_VERSION_HEADER));
Self {
request_ip_address,
avt_cookie_token,
session_token,
session_user_token,
trace_id,
http_host,
http_origin,
http_referer,
http_user_agent,
artcraft_version,
}
}
}
/// Read a header as a UTF-8 string; `None` if absent or not valid UTF-8.
fn get_header_value(request: &HttpRequest, header_name: HeaderName) -> Option<String> {
request.headers()
.get(header_name)
.and_then(|value| value.to_str().ok())
.map(|value| value.to_string())
}
@@ -9,7 +9,6 @@ use crate::http_server::deprecated_endpoints::categories::get_category::get_cate
use crate::http_server::deprecated_endpoints::categories::tts::assign_tts_category::assign_tts_category_handler;
use crate::http_server::deprecated_endpoints::categories::tts::list_fully_computed_assigned_tts_categories::list_fully_computed_assigned_tts_categories::list_fully_computed_assigned_tts_categories_handler;
use crate::http_server::deprecated_endpoints::categories::tts::list_tts_categories::list_tts_categories_handler;
use crate::http_server::deprecated_endpoints::categories::tts::list_tts_model_assigned_categories::list_tts_model_assigned_categories_handler;
use crate::http_server::deprecated_endpoints::conversion::enqueue_fbx_to_gltf_handler::enqueue_fbx_to_gltf_handler;
use crate::http_server::deprecated_endpoints::conversion::enqueue_render_engine_scene_to_video_handler::enqueue_render_engine_scene_to_video_handler;
use crate::http_server::deprecated_endpoints::engine::create_scene_handler::create_scene_handler;
@@ -18,7 +17,6 @@ use crate::http_server::deprecated_endpoints::engine::update_scene_handler::upda
use crate::http_server::deprecated_endpoints::events::list_events::list_events_handler;
use crate::http_server::deprecated_endpoints::flags::design_refresh_flag::disable_design_refresh_flag_handler::disable_design_refresh_flag_handler;
use crate::http_server::deprecated_endpoints::flags::design_refresh_flag::enable_design_refresh_flag_handler::enable_design_refresh_flag_handler;
use crate::http_server::deprecated_endpoints::leaderboard::get_leaderboard::leaderboard_handler;
use crate::http_server::deprecated_endpoints::media_uploads::list_user_media_uploads_of_type::list_user_media_uploads_of_type_handler;
use crate::http_server::deprecated_endpoints::media_uploads::upload_audio::upload_audio_handler;
use crate::http_server::deprecated_endpoints::media_uploads::upload_image::upload_image_handler;
@@ -94,11 +92,6 @@ where
.add_post("/v1/conversion/enqueue_render_engine_scene", enqueue_render_engine_scene_to_video_handler)
.into_app();
let app = app.service(web::resource("/leaderboard")
.route(web::get().to(leaderboard_handler))
.route(web::head().to(|| HttpResponse::Ok()))
);
// ==================== FakeYou Frontend Application State ====================
let mut app = RouteBuilder::from_app(app)
@@ -107,6 +100,8 @@ where
// ==================== Stats ====================
// NB: Deprecated; serves a hardcoded snapshot for legacy clients that still poll it.
#[allow(deprecated)]
let mut app = RouteBuilder::from_app(app)
.add_get("/v1/stats/queues", get_unified_queue_stats_handler)
.into_app();
@@ -237,14 +232,6 @@ where
.route(web::head().to(|| HttpResponse::Ok()))
)
)
.service(
web::scope("/assignments")
.service(
web::resource("/tts/{token}")
.route(web::get().to(list_tts_model_assigned_categories_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
)
)
}
@@ -364,32 +351,6 @@ where
)
}
//fn add_image_gen_routes<T,B> (app:App<T>)-> App<T>
// where
// B: MessageBody,
// T: ServiceFactory<
// ServiceRequest,
// Config = (),
// Response = ServiceResponse<B>,
// Error = Error,
// InitError = (),
// >,
//{
// //
// app.service(
// web::scope("/v1/image_gen")
// .service(
// web::scope("/upload")
// .route("/lora", web::post().to(enqueue_image_generation_request))
// .route("/model", web::post().to(enqueue_image_generation_request))
// )
// .service(
// web::scope("/enqueue")
// .route("/inference", web::post().to(enqueue_image_generation_request))
// )
// )
//}
// ==================== Engine Routes ====================
fn add_engine_routes<T, B>(app: App<T>) -> App<T>
@@ -5,16 +5,12 @@ use actix_web::error::Error;
use actix_web::{web, App, HttpResponse};
use crate::http_server::endpoints::tts::delete_tts_model::delete_tts_model_handler;
use crate::http_server::endpoints::tts::delete_tts_result::delete_tts_inference_result_handler;
use crate::http_server::endpoints::tts::edit_tts_model::edit_tts_model_handler;
use crate::http_server::endpoints::tts::edit_tts_result::edit_tts_inference_result_handler;
use crate::http_server::endpoints::tts::enqueue_infer_tts_handler::enqueue_infer_tts_handler::enqueue_infer_tts_handler;
use crate::http_server::endpoints::tts::enqueue_upload_tts_model::upload_tts_model_handler;
use crate::http_server::endpoints::tts::get_pending_tts_inference_job_count::get_pending_tts_inference_job_count_handler;
use crate::http_server::endpoints::tts::get_tts_inference_job_status::get_tts_inference_job_status_handler;
use crate::http_server::endpoints::tts::get_tts_inference_job_status_handler::get_tts_inference_job_status_handler;
use crate::http_server::endpoints::tts::get_tts_model::get_tts_model_handler;
use crate::http_server::endpoints::tts::get_tts_model_use_count::get_tts_model_use_count_handler;
use crate::http_server::endpoints::tts::get_tts_result::get_tts_inference_result_handler;
use crate::http_server::endpoints::tts::get_tts_upload_model_job_status::get_tts_upload_model_job_status_handler;
use crate::http_server::endpoints::tts::list_tts_models::list_tts_models_handler;
use crate::http_server::endpoints::tts::search_tts_models_handler::search_tts_models_handler;
@@ -40,11 +36,6 @@ where
// NB(bt,2024-04-03): Older root /tts/* routes should be drained of traffic. They are used in the Twitch Streamer API though.
.service(
web::scope("/tts")
.service(
web::resource("/upload")
.route(web::post().to(upload_tts_model_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/inference")
.route(web::post().to(enqueue_infer_tts_handler))
@@ -70,31 +61,16 @@ where
.route(web::post().to(delete_tts_model_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/model/{model_token}/count")
.route(web::get().to(get_tts_model_use_count_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/model/{model_token}/edit")
.route(web::post().to(edit_tts_model_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/result/{token}")
.route(web::get().to(get_tts_inference_result_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/result/{token}/edit")
.route(web::post().to(edit_tts_inference_result_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/result/{token}/delete")
.route(web::post().to(delete_tts_inference_result_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(
web::resource("/job/{token}")
.route(web::get().to(get_tts_inference_job_status_handler))
@@ -0,0 +1 @@
pub mod require_api_or_web_session;
@@ -0,0 +1,190 @@
use std::marker::PhantomData;
use actix_web::http::header;
use actix_web::HttpRequest;
use log::warn;
use sqlx::{Executor, MySql};
use actix_artcraft::sessions::anonymous_visitor_tracking::avt_cookie_manager::AvtCookieManager;
use mysql_queries::queries::users::api_or_web_sessions::api_or_web_session_user_record::ApiOrWebSessionUserRecord;
use mysql_queries::queries::users::api_or_web_sessions::get_api_or_web_session_user_by_api_key::{
get_api_or_web_session_user_by_api_key, GetApiOrWebSessionUserByApiKeyArgs,
};
use mysql_queries::queries::users::api_or_web_sessions::get_api_or_web_session_user_by_session_token::{
get_api_or_web_session_user_by_session_token, GetApiOrWebSessionUserBySessionTokenArgs,
};
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
use tokens::tokens::api_keys::ApiKeyToken;
use tokens::tokens::users::UserToken;
use crate::http_server::common_responses::common_web_error::CommonWebError;
use crate::http_server::user_lookup::user_session::session_utils::session_checker::SessionChecker;
use crate::http_server::web_utils::get_authorization_header_api_key::get_authorization_header_api_key;
const ADMIN_ROLE_SLUG: &str = "admin";
/// The authenticated identity behind a request that may come from either an API client
/// (`Authorization` header API key) or the website (session cookie). Fields common to both auth
/// paths are populated uniformly, so handlers don't need to care which path authenticated the
/// user.
pub struct ApiOrWebSession {
pub session_type: SessionType,
pub user_token: UserToken,
pub username: String,
pub display_name: String,
pub email_address: String,
pub user_role_slug: String,
pub can_ban_users: bool,
/// Only present for API-key sessions.
pub maybe_api_key_token: Option<ApiKeyToken>,
/// Anonymous visitor tracking cookie. Only ever present for web sessions.
pub maybe_avt_token: Option<AnonymousVisitorTrackingToken>,
}
/// Which auth path authenticated the request.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum SessionType {
Api,
WebSession,
}
impl ApiOrWebSession {
pub fn is_mod(&self) -> bool {
self.can_ban_users || self.user_role_slug == ADMIN_ROLE_SLUG
}
}
/// Authenticate a request as either an API-key user or a web-session (cookie) user.
///
/// If an `Authorization` header is present, ONLY the API-key lookup is attempted — a malformed or
/// unknown key is a 401, never a fallthrough to cookies. Otherwise the session cookie is decoded
/// and looked up, and the AVT (anonymous visitor tracking) cookie is captured alongside it.
///
/// Neither path is cached: every request performs a fresh MySQL lookup. `mysql_executor` can be
/// any sqlx executor — prefer passing an already-open connection (`&mut *connection`) so the
/// lookup reuses it rather than acquiring a fresh one from the pool.
pub async fn require_api_or_web_session<'c, E>(
http_request: &HttpRequest,
session_checker: &SessionChecker,
avt_cookie_manager: &AvtCookieManager,
mysql_executor: E,
) -> Result<ApiOrWebSession, CommonWebError>
where E: 'c + Executor<'c, Database = MySql>
{
if http_request.headers().contains_key(header::AUTHORIZATION) {
authenticate_by_api_key(http_request, mysql_executor).await
} else {
authenticate_by_web_session(http_request, session_checker, avt_cookie_manager, mysql_executor).await
}
}
async fn authenticate_by_api_key<'c, E>(
http_request: &HttpRequest,
mysql_executor: E,
) -> Result<ApiOrWebSession, CommonWebError>
where E: 'c + Executor<'c, Database = MySql>
{
let api_key = match get_authorization_header_api_key(http_request) {
Some(api_key) => api_key,
None => {
warn!("Authorization header present but not a usable API key");
return Err(CommonWebError::NotAuthorized);
}
};
let maybe_record = get_api_or_web_session_user_by_api_key(GetApiOrWebSessionUserByApiKeyArgs {
api_key: &api_key,
mysql_executor,
phantom: PhantomData,
})
.await
.map_err(|err| {
warn!("API key user lookup error: {:?}", err);
CommonWebError::from_error(err)
})?;
// A missing or soft-deleted key (or a key whose owner no longer exists) is a 401, not a leak of
// which case occurred.
let record = match maybe_record {
Some(record) => record,
None => {
warn!("No live API key user for presented key: {:?}", api_key);
return Err(CommonWebError::NotAuthorized);
}
};
into_session_rejecting_banned(record, SessionType::Api, None)
}
async fn authenticate_by_web_session<'c, E>(
http_request: &HttpRequest,
session_checker: &SessionChecker,
avt_cookie_manager: &AvtCookieManager,
mysql_executor: E,
) -> Result<ApiOrWebSession, CommonWebError>
where E: 'c + Executor<'c, Database = MySql>
{
let maybe_session_token = session_checker
.get_session_token(http_request)
.map_err(|err| {
warn!("Session cookie decode error: {:?}", err);
CommonWebError::from_error(err)
})?;
let session_token = match maybe_session_token {
Some(session_token) => session_token,
None => {
warn!("not logged in");
return Err(CommonWebError::NotAuthorized);
}
};
let maybe_record = get_api_or_web_session_user_by_session_token(GetApiOrWebSessionUserBySessionTokenArgs {
session_token: &session_token,
mysql_executor,
phantom: PhantomData,
})
.await
.map_err(|err| {
warn!("Web session user lookup error: {:?}", err);
CommonWebError::from_error(err)
})?;
let record = match maybe_record {
Some(record) => record,
None => {
warn!("Valid cookie; invalid session: {}", session_token);
return Err(CommonWebError::NotAuthorized);
}
};
let maybe_avt_token = avt_cookie_manager.get_avt_token_from_request(http_request);
into_session_rejecting_banned(record, SessionType::WebSession, maybe_avt_token)
}
fn into_session_rejecting_banned(
record: ApiOrWebSessionUserRecord,
session_type: SessionType,
maybe_avt_token: Option<AnonymousVisitorTrackingToken>,
) -> Result<ApiOrWebSession, CommonWebError> {
if record.is_banned {
warn!("user is banned: {:?}", record.user_token.as_str());
return Err(CommonWebError::NotAuthorized);
}
Ok(ApiOrWebSession {
session_type,
user_token: record.user_token,
username: record.username,
display_name: record.display_name,
email_address: record.email_address,
user_role_slug: record.user_role_slug,
can_ban_users: record.can_ban_users,
maybe_api_key_token: record.maybe_api_key_token,
maybe_avt_token,
})
}
@@ -1,2 +1,3 @@
pub mod api_keys;
pub mod api_or_web_session;
pub mod user_session;
@@ -151,10 +151,6 @@ pub async fn setup_dependencies(server_hostname: &str) -> AnyhowResult<SetupResu
easyenv::get_env_duration_seconds_or_default("TTS_MODEL_CATEGORY_ASSIGNMENTS_CACHE_TTL_SECONDS", Duration::from_secs(60))
);
let leaderboard_cache = SingleItemTtlCache::create_with_duration(
easyenv::get_env_duration_seconds_or_default("LEADERBOARD_CACHE_TTL_SECONDS", Duration::from_secs(60))
);
let inference_queue_length_cache = SingleItemTtlCache::create_with_duration(
easyenv::get_env_duration_seconds_or_default("INFERENCE_QUEUE_LENGTH_CACHE_TTL_SECONDS", Duration::from_secs(30))
);
@@ -251,12 +247,7 @@ pub async fn setup_dependencies(server_hostname: &str) -> AnyhowResult<SetupResu
tts_model_list: voice_list_cache,
database_tts_category_list: database_tts_category_list_cache,
tts_model_category_assignments: tts_model_category_assignments_cache,
leaderboard: leaderboard_cache,
inference_queue_length: inference_queue_length_cache,
queue_stats: SingleItemTtlCache::create_with_duration(
easyenv::get_env_duration_seconds_or_default(
"QUEUE_STATS_CACHE_TTL_SECONDS",
Duration::from_secs(60))),
featured_media_files_sieve: ArcTtlSieve::with_capacity_and_ttl_duration(
easyenv::get_env_num("FEATURED_MEDIA_FILES_CACHE_SIZE", 25)?,
easyenv::get_env_duration_seconds_or_default("FEATURED_MEDIA_FILES_TTL_SECONDS", Duration::from_secs(60)),
@@ -11,12 +11,10 @@ pub fn setup_static_feature_flags(paging_flags: PagingFlags) -> AnyhowResult<Sta
// Permanent (control plane / safety) flags : disabling features
global_429_pushback_filter_enabled: easyenv::get_env_bool_or_default("FF_GLOBAL_429_PUSHBACK_FILTER_ENABLED", false),
disable_unified_queue_stats_endpoint: easyenv::get_env_bool_or_default("FF_DISABLE_QUEUE_STATS_ENDPOINT", false),
disable_inference_queue_length_endpoint: easyenv::get_env_bool_or_default("FF_DISABLE_INFERENCE_QUEUE_LENGTH_ENDPOINT", false),
disable_tts_model_list_endpoint: easyenv::get_env_bool_or_default("FF_DISABLE_TTS_MODEL_LIST_ENDPOINT", false),
// Refresh rates
frontend_unified_queue_stats_refresh_interval_millis: easyenv::get_env_num("FF_FRONTEND_QUEUE_STATS_REFRESH_INTERVAL_MILLIS", 15_000)?,
frontend_pending_inference_refresh_interval_millis: easyenv::get_env_num("FF_FRONTEND_PENDING_INFERENCE_REFRESH_INTERVAL_MILLIS", 15_000)?,
// Bans
@@ -1,9 +1,7 @@
use crate::configs::app_startup::username_set::UsernameSet;
use crate::configs::static_api_tokens::StaticApiTokenSet;
use crate::http_server::deprecated_endpoints::categories::tts::list_fully_computed_assigned_tts_categories::list_fully_computed_assigned_tts_categories::ModelTokensByCategoryToken;
use crate::http_server::deprecated_endpoints::leaderboard::get_leaderboard::LeaderboardInfo;
use crate::http_server::endpoints::media_files::list::list_featured_media_files_handler::ListFeaturedMediaFilesQueryParams;
use crate::http_server::endpoints::stats::result_transformer::CacheableQueueStats;
use crate::http_server::endpoints::tts::list_tts_models::TtsModelRecordForResponse;
use crate::http_server::user_lookup::user_session::session_utils::session_checker::SessionChecker;
use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter;
@@ -197,18 +195,11 @@ pub struct EphemeralInMemoryCaches {
/// This is approximately O(n^3) and recursively generates all super-category membership.
pub tts_model_category_assignments: SingleItemTtlCache<ModelTokensByCategoryToken>,
/// Stats on generic inference queue and legacy TTS queue (combined).
/// The frontend will consult a distributed cache and use the monotonic DB time as a
/// vector clock.
pub queue_stats: SingleItemTtlCache<CacheableQueueStats>,
/// Generic inference queue length
/// The frontend will consult a distributed cache and use the monotonic DB time as a
/// vector clock.
pub inference_queue_length: SingleItemTtlCache<InferenceQueueLengthResult>,
pub leaderboard: SingleItemTtlCache<LeaderboardInfo>,
/// Cache of featured media files
pub featured_media_files_sieve: ArcTtlSieve<ListFeaturedMediaFilesQueryParams, FeaturedMediaFileListPage>,
}
@@ -235,19 +226,12 @@ pub struct StaticFeatureFlags {
/// If we're suffering an outage, set custom text for the alert message.
pub maybe_status_alert_custom_message: Option<String>,
/// Disable the live `/v1/stats/queues` endpoint for all users and serve a static value instead.
pub disable_unified_queue_stats_endpoint: bool,
/// Disable the live `/v1/model_inference/queue_length` endpoint for all users and serve a static value instead.
pub disable_inference_queue_length_endpoint: bool,
/// Disable the live `/tts/list` endpoint for all users and serve a static value instead.
pub disable_tts_model_list_endpoint: bool,
/// Tell the frontend client how fast to refresh their view of queue stats.
/// During an attack, we may want this to go extremely slow.
pub frontend_unified_queue_stats_refresh_interval_millis: u64,
/// Tell the frontend client how fast to refresh their view of the pending inference count.
/// During an attack, we may want this to go extremely slow.
pub frontend_pending_inference_refresh_interval_millis: u64,