diff --git a/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4_edit_image_webhook.rs b/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4_edit_image_webhook.rs index 3c56f13cba..58689a8efb 100644 --- a/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4_edit_image_webhook.rs +++ b/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4_edit_image_webhook.rs @@ -1,6 +1,7 @@ use crate::creds::fal_api_key::FalApiKey; use crate::error::classify_fal_error::classify_fal_error; use crate::error::fal_error_plus::FalErrorPlus; +use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents}; use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4_edit_image::{bytedance_seedream_v4_edit_image, BytedanceSeedreamV4EditImageInput}; use fal::webhook::WebhookResponse; use reqwest::IntoUrl; @@ -53,6 +54,23 @@ pub enum EnqueueBytedanceSeedreamV4EditImageSize { Auto4k, } + +impl FalRequestCostCalculator for EnqueueBytedanceSeedreamV4EditImageArgs<'_, U> { + fn calculate_cost_in_cents(&self) -> UsdCents { + // Your request will cost $0.03 per image. + let unit_cost = 3; + let cost = match self.num_images { + None => unit_cost, + Some(EnqueueBytedanceSeedreamV4EditImageNumImages::One) => unit_cost, + Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Two) => unit_cost * 2, + Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Three) => unit_cost * 3, + Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Four) => unit_cost * 4, + }; + cost as UsdCents + } +} + + pub async fn enqueue_bytedance_seedream_v4_edit_image_webhook( args: EnqueueBytedanceSeedreamV4EditImageArgs<'_, R> ) -> Result { diff --git a/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4p5_edit_image_webhook.rs b/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4p5_edit_image_webhook.rs index 558971422f..95b543512f 100644 --- a/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4p5_edit_image_webhook.rs +++ b/crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4p5_edit_image_webhook.rs @@ -1,6 +1,7 @@ use crate::creds::fal_api_key::FalApiKey; use crate::error::classify_fal_error::classify_fal_error; use crate::error::fal_error_plus::FalErrorPlus; +use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents}; use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4p5_edit_image::{bytedance_seedream_v4p5_edit_image, BytedanceSeedreamV4p5EditImageInput}; use fal::webhook::WebhookResponse; use reqwest::IntoUrl; @@ -52,6 +53,23 @@ pub enum EnqueueBytedanceSeedreamV4p5EditImageSize { Auto4k, } + +impl FalRequestCostCalculator for EnqueueBytedanceSeedreamV4p5EditImageArgs<'_, U> { + fn calculate_cost_in_cents(&self) -> UsdCents { + // Your request will cost $0.04 per image. + let unit_cost = 4; + let cost = match self.num_images { + None => unit_cost, + Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::One) => unit_cost, + Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Two) => unit_cost * 2, + Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Three) => unit_cost * 3, + Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Four) => unit_cost * 4, + }; + cost as UsdCents + } +} + + pub async fn enqueue_bytedance_seedream_v4p5_edit_image_webhook( args: EnqueueBytedanceSeedreamV4p5EditImageArgs<'_, R> ) -> Result { diff --git a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4_text_to_image_webhook.rs b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4_text_to_image_webhook.rs index 92a2c427fe..c859f4e7b4 100644 --- a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4_text_to_image_webhook.rs +++ b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4_text_to_image_webhook.rs @@ -1,6 +1,7 @@ use crate::creds::fal_api_key::FalApiKey; use crate::error::classify_fal_error::classify_fal_error; use crate::error::fal_error_plus::FalErrorPlus; +use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents}; use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4_text_to_image::BytedanceSeedreamV4TextToImageInput; use fal::prelude::fal_ai::bytedance::seedream::bytedance_seedream_v4_text_to_image::bytedance_seedream_v4_text_to_image; use fal::webhook::WebhookResponse; @@ -53,6 +54,24 @@ pub enum EnqueueBytedanceSeedreamV4TextToImageSize { Auto4k, } + +impl FalRequestCostCalculator for EnqueueBytedanceSeedreamV4TextToImageArgs<'_, U> { + fn calculate_cost_in_cents(&self) -> UsdCents { + // Copied from edit image version + // Your request will cost $0.03 per image. + let unit_cost = 3; + let cost = match self.num_images { + None => unit_cost, + Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::One) => unit_cost, + Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Two) => unit_cost * 2, + Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Three) => unit_cost * 3, + Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Four) => unit_cost * 4, + }; + cost as UsdCents + } +} + + pub async fn enqueue_bytedance_seedream_v4_text_to_image_webhook( args: EnqueueBytedanceSeedreamV4TextToImageArgs<'_, R> ) -> Result { diff --git a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4p5_text_to_image_webhook.rs b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4p5_text_to_image_webhook.rs index b47b2133bc..a5cf2f6c3c 100644 --- a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4p5_text_to_image_webhook.rs +++ b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4p5_text_to_image_webhook.rs @@ -4,6 +4,8 @@ use crate::error::fal_error_plus::FalErrorPlus; use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4p5_text_to_image::{bytedance_seedream_v4p5_text_to_image, BytedanceSeedreamV4p5TextToImageInput}; use fal::webhook::WebhookResponse; use reqwest::IntoUrl; +use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents}; +use crate::requests::webhook::image::edit::enqueue_bytedance_seedream_v4p5_edit_image_webhook::{EnqueueBytedanceSeedreamV4p5EditImageArgs, EnqueueBytedanceSeedreamV4p5EditImageNumImages}; pub struct EnqueueBytedanceSeedreamV4p5TextToImageArgs<'a, R: IntoUrl> { // Request required @@ -51,6 +53,24 @@ pub enum EnqueueBytedanceSeedreamV4p5TextToImageSize { Auto4k, } + +impl FalRequestCostCalculator for EnqueueBytedanceSeedreamV4p5TextToImageArgs<'_, U> { + fn calculate_cost_in_cents(&self) -> UsdCents { + // NB: Copied from edit image costs + // Your request will cost $0.04 per image. + let unit_cost = 4; + let cost = match self.num_images { + None => unit_cost, + Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::One) => unit_cost, + Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Two) => unit_cost * 2, + Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Three) => unit_cost * 3, + Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Four) => unit_cost * 4, + }; + cost as UsdCents + } +} + + pub async fn enqueue_bytedance_seedream_v4p5_text_to_image_webhook( args: EnqueueBytedanceSeedreamV4p5TextToImageArgs<'_, R> ) -> Result { diff --git a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_gemini_25_flash_text_to_image_webhook.rs b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_gemini_25_flash_text_to_image_webhook.rs index 6ae117ec6c..dd38395f04 100644 --- a/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_gemini_25_flash_text_to_image_webhook.rs +++ b/crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_gemini_25_flash_text_to_image_webhook.rs @@ -5,6 +5,8 @@ use fal::prelude::fal_ai::gemini_25_flash_image::text_to_image::gemini_25_flash_ use fal::prelude::fal_ai::gemini_25_flash_image::text_to_image::Gemini25FlashTextToImageInput; use fal::webhook::WebhookResponse; use reqwest::IntoUrl; +use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents}; +use crate::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::{Gemini25FlashEditArgs, Gemini25FlashEditNumImages}; pub struct Gemini25FlashTextToImageArgs<'a, R: IntoUrl> { // Request required @@ -46,6 +48,21 @@ pub enum Gemini25FlashTextToImageAspectRatio { NineBySixteen, // NB: No NineByTwentyOne ? } + +impl FalRequestCostCalculator for Gemini25FlashTextToImageArgs<'_, U> { + fn calculate_cost_in_cents(&self) -> UsdCents { + // NB: This was copied from the image-to-image case. + // Your request will cost $0.039 per image. For $1.00, you can run this model 25 times. + match self.num_images { + Gemini25FlashTextToImageNumImages::One => 4, + Gemini25FlashTextToImageNumImages::Two => 8, + Gemini25FlashTextToImageNumImages::Three => 12, + Gemini25FlashTextToImageNumImages::Four => 16, + } + } +} + + pub async fn enqueue_gemini_25_flash_text_to_image_webhook( args: Gemini25FlashTextToImageArgs<'_, R> ) -> Result { diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4_multi_function_image_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4_multi_function_image_gen_handler.rs index 57ad936884..7cca9d062c 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4_multi_function_image_gen_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4_multi_function_image_gen_handler.rs @@ -2,6 +2,7 @@ use std::collections::HashSet; use std::iter::FromIterator; use std::sync::Arc; +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::common_responses::media::media_links_builder::MediaLinksBuilder; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; @@ -19,6 +20,7 @@ use enums::common::generation_provider::GenerationProvider; use enums::common::model_type::ModelType; use enums::common::visibility::Visibility; use fal_client::creds::open_ai_api_key::OpenAiApiKey; +use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator; use fal_client::requests::webhook::image::edit::enqueue_bytedance_seedream_v4_edit_image_webhook::{enqueue_bytedance_seedream_v4_edit_image_webhook, EnqueueBytedanceSeedreamV4EditImageArgs, EnqueueBytedanceSeedreamV4EditImageNumImages, EnqueueBytedanceSeedreamV4EditImageSize}; use fal_client::requests::webhook::image::text::enqueue_bytedance_seedream_v4_text_to_image_webhook::{enqueue_bytedance_seedream_v4_text_to_image_webhook, EnqueueBytedanceSeedreamV4TextToImageArgs, EnqueueBytedanceSeedreamV4TextToImageNumImages, EnqueueBytedanceSeedreamV4TextToImageSize}; use http_server_common::request::get_request_ip::get_request_ip; @@ -26,6 +28,7 @@ use log::{error, info, warn}; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs}; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection}; use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; @@ -33,6 +36,7 @@ use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPrompt use server_environment::ServerEnvironment; use sqlx::pool::PoolConnection; use sqlx::{Acquire, MySql}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; use utoipa::ToSchema; @@ -73,15 +77,12 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( .avt_cookie_manager .get_avt_token_from_request(&http_request); - // TODO: Limit usage for new accounts. Billing, free credits metering, etc. - - //let user_session = match maybe_user_session { - // Some(session) => session, - // None => { - // warn!("not logged in"); - // return Err(CommonWebError::NotAuthorized); - // } - //}; + let user_token = match maybe_user_session.as_ref() { + Some(session) => &session.user_token, + None => { + return Err(CommonWebError::NotAuthorized); + } + }; if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { return Err(CommonWebError::BadInputWithSimpleMessage(reason)); @@ -109,6 +110,8 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( info!("Fal webhook URL: {}", server_state.fal.webhook_url); + let apriori_job_token = InferenceJobToken::generate(); + let fal_result; if let Some(input_image_urls) = image_urls.as_deref() { @@ -149,6 +152,17 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_bytedance_seedream_v4_edit_image_webhook(args) .await .map_err(|err| { @@ -193,6 +207,17 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_bytedance_seedream_v4_text_to_image_webhook(args) .await .map_err(|err| { @@ -223,9 +248,7 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( let prompt_result = insert_prompt(InsertPromptArgs { maybe_apriori_prompt_token: None, prompt_type: PromptType::ArtcraftApp, - maybe_creator_user_token: maybe_user_session - .as_ref() - .map(|s| &s.user_token), + maybe_creator_user_token: Some(&user_token), maybe_model_type: Some(ModelType::Seedream4), maybe_generation_provider: Some(GenerationProvider::Artcraft), maybe_positive_prompt: request.prompt.as_deref(), @@ -264,13 +287,14 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler( } } - let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs { + let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs { + apriori_job_token: &apriori_job_token, uuid_idempotency_token: &request.uuid_idempotency_token, maybe_external_third_party_id: &external_job_id, fal_category: FalCategory::ImageGeneration, maybe_inference_args: None, maybe_prompt_token: prompt_token.as_ref(), - maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token), + maybe_creator_user_token: Some(&user_token), maybe_avt_token: maybe_avt_token.as_ref(), creator_ip_address: &ip_address, creator_set_visibility: Visibility::Public, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4p5_multi_function_image_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4p5_multi_function_image_gen_handler.rs index e52a77c4dc..cc684e6722 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4p5_multi_function_image_gen_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4p5_multi_function_image_gen_handler.rs @@ -2,6 +2,7 @@ use std::collections::HashSet; use std::iter::FromIterator; use std::sync::Arc; +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::common_responses::media::media_links_builder::MediaLinksBuilder; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; @@ -19,6 +20,7 @@ use enums::common::generation_provider::GenerationProvider; use enums::common::model_type::ModelType; use enums::common::visibility::Visibility; use fal_client::creds::open_ai_api_key::OpenAiApiKey; +use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator; use fal_client::requests::webhook::image::edit::enqueue_bytedance_seedream_v4p5_edit_image_webhook::{enqueue_bytedance_seedream_v4p5_edit_image_webhook, EnqueueBytedanceSeedreamV4p5EditImageArgs, EnqueueBytedanceSeedreamV4p5EditImageNumImages, EnqueueBytedanceSeedreamV4p5EditImageSize}; use fal_client::requests::webhook::image::text::enqueue_bytedance_seedream_v4p5_text_to_image_webhook::{enqueue_bytedance_seedream_v4p5_text_to_image_webhook, EnqueueBytedanceSeedreamV4p5TextToImageArgs, EnqueueBytedanceSeedreamV4p5TextToImageNumImages, EnqueueBytedanceSeedreamV4p5TextToImageSize}; use http_server_common::request::get_request_ip::get_request_ip; @@ -26,6 +28,7 @@ use log::{error, info, warn}; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs}; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection}; use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; @@ -33,6 +36,7 @@ use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPrompt use server_environment::ServerEnvironment; use sqlx::pool::PoolConnection; use sqlx::{Acquire, MySql}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; use utoipa::ToSchema; @@ -73,15 +77,12 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( .avt_cookie_manager .get_avt_token_from_request(&http_request); - // TODO: Limit usage for new accounts. Billing, free credits metering, etc. - - //let user_session = match maybe_user_session { - // Some(session) => session, - // None => { - // warn!("not logged in"); - // return Err(CommonWebError::NotAuthorized); - // } - //}; + let user_token = match maybe_user_session.as_ref() { + Some(session) => &session.user_token, + None => { + return Err(CommonWebError::NotAuthorized); + } + }; if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { return Err(CommonWebError::BadInputWithSimpleMessage(reason)); @@ -109,6 +110,8 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( info!("Fal webhook URL: {}", server_state.fal.webhook_url); + let apriori_job_token = InferenceJobToken::generate(); + let fal_result; if let Some(input_image_urls) = image_urls.as_deref() { @@ -148,6 +151,17 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_bytedance_seedream_v4p5_edit_image_webhook(args) .await .map_err(|err| { @@ -191,6 +205,17 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_bytedance_seedream_v4p5_text_to_image_webhook(args) .await .map_err(|err| { @@ -221,9 +246,7 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( let prompt_result = insert_prompt(InsertPromptArgs { maybe_apriori_prompt_token: None, prompt_type: PromptType::ArtcraftApp, - maybe_creator_user_token: maybe_user_session - .as_ref() - .map(|s| &s.user_token), + maybe_creator_user_token: Some(&user_token), maybe_model_type: Some(ModelType::Seedream4p5), maybe_generation_provider: Some(GenerationProvider::Artcraft), maybe_positive_prompt: request.prompt.as_deref(), @@ -262,13 +285,14 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler( } } - let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs { + let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs { + apriori_job_token: &apriori_job_token, uuid_idempotency_token: &request.uuid_idempotency_token, maybe_external_third_party_id: &external_job_id, fal_category: FalCategory::ImageGeneration, maybe_inference_args: None, maybe_prompt_token: prompt_token.as_ref(), - maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token), + maybe_creator_user_token: Some(&user_token), maybe_avt_token: maybe_avt_token.as_ref(), creator_ip_address: &ip_address, creator_set_visibility: Visibility::Public, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_multi_function_image_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_multi_function_image_gen_handler.rs index cd192f84af..fa1c0b1a42 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_multi_function_image_gen_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_multi_function_image_gen_handler.rs @@ -2,6 +2,7 @@ use std::collections::HashSet; use std::iter::FromIterator; use std::sync::Arc; +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::common_responses::media::media_links_builder::MediaLinksBuilder; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; @@ -22,6 +23,7 @@ use enums::common::payments_namespace::PaymentsNamespace; use enums::common::stripe_subscription_status::StripeSubscriptionStatus; use enums::common::visibility::Visibility; use fal_client::creds::open_ai_api_key::OpenAiApiKey; +use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator; use fal_client::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::{enqueue_gemini_25_flash_edit_webhook, Gemini25FlashEditArgs, Gemini25FlashEditAspectRatio, Gemini25FlashEditNumImages}; use fal_client::requests::webhook::image::edit::enqueue_nano_banana_pro_edit_image_webhook::{enqueue_nano_banana_pro_image_edit_webhook, EnqueueNanoBananaProEditImageArgs, EnqueueNanoBananaProEditImageAspectRatio, EnqueueNanoBananaProEditImageNumImages, EnqueueNanoBananaProEditImageResolution}; use fal_client::requests::webhook::image::text::enqueue_gemini_25_flash_text_to_image_webhook::{enqueue_gemini_25_flash_text_to_image_webhook, Gemini25FlashTextToImageArgs, Gemini25FlashTextToImageAspectRatio, Gemini25FlashTextToImageNumImages}; @@ -31,6 +33,7 @@ use log::{error, info, warn}; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory; use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs}; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection}; use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; @@ -39,6 +42,7 @@ use mysql_queries::queries::users::user_subscriptions::find_subscription_for_own use server_environment::ServerEnvironment; use sqlx::pool::PoolConnection; use sqlx::{Acquire, MySql}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; use utoipa::ToSchema; @@ -83,39 +87,31 @@ pub async fn nano_banana_multi_function_image_gen_handler( .avt_cookie_manager .get_avt_token_from_request(&http_request); - let maybe_user_token = maybe_user_session - .as_ref() - .map(|s| &s.user_token); - - // TODO: Limit usage for new accounts. Billing, free credits metering, etc. - - //let user_session = match maybe_user_session { - // Some(session) => session, - // None => { - // warn!("not logged in"); - // return Err(CommonWebError::NotAuthorized); - // } - //}; + let user_token = match maybe_user_session.as_ref() { + Some(session) => &session.user_token, + None => { + return Err(CommonWebError::NotAuthorized); + } + }; let mut downgrade_for_free_user = true; - if let Some(user_token) = maybe_user_token.as_ref() { - let result = find_subscription_for_owner_user_using_connection( - user_token, - PaymentsNamespace::Artcraft, - &mut mysql_connection, - ).await; + let result = find_subscription_for_owner_user_using_connection( + user_token, + PaymentsNamespace::Artcraft, + &mut mysql_connection, + ).await; - if let Ok(Some(subscription)) = result { - info!("User {:?} has subscription: {:?} (stripe customer: {:?}, status: {:?})", - user_token, - subscription.token, - subscription.stripe_customer_id, - subscription.stripe_subscription_status); - // NB: Failing open means subscribers might get fewer results, but they're free right now. - if subscription.stripe_subscription_status == StripeSubscriptionStatus::Active { - downgrade_for_free_user = false; - } + if let Ok(Some(subscription)) = result { + info!("User {:?} has subscription: {:?} (stripe customer: {:?}, status: {:?})", + user_token, + subscription.token, + subscription.stripe_customer_id, + subscription.stripe_subscription_status); + + // NB: Failing open means subscribers might get fewer results, but they're free right now. + if subscription.stripe_subscription_status == StripeSubscriptionStatus::Active { + downgrade_for_free_user = false; } } @@ -143,6 +139,8 @@ pub async fn nano_banana_multi_function_image_gen_handler( info!("Fal webhook URL: {}", server_state.fal.webhook_url); + let apriori_job_token = InferenceJobToken::generate(); + let fal_result; if let Some(input_image_urls) = image_urls.as_deref() { @@ -184,6 +182,17 @@ pub async fn nano_banana_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_gemini_25_flash_edit_webhook(args) .await .map_err(|err| { @@ -229,6 +238,17 @@ pub async fn nano_banana_multi_function_image_gen_handler( api_key: &server_state.fal.api_key, }; + let cost = args.calculate_cost_in_cents(); + + info!("Charging wallet: {}", cost); + + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + &mut mysql_connection, + ).await?; + fal_result = enqueue_gemini_25_flash_text_to_image_webhook(args) .await .map_err(|err| { @@ -259,7 +279,7 @@ pub async fn nano_banana_multi_function_image_gen_handler( let prompt_result = insert_prompt(InsertPromptArgs { maybe_apriori_prompt_token: None, prompt_type: PromptType::ArtcraftApp, - maybe_creator_user_token: maybe_user_token, + maybe_creator_user_token: Some(&user_token), maybe_model_type: Some(ModelType::NanoBanana), maybe_generation_provider: Some(GenerationProvider::Artcraft), maybe_positive_prompt: request.prompt.as_deref(), @@ -298,13 +318,14 @@ pub async fn nano_banana_multi_function_image_gen_handler( } } - let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs { + let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs { + apriori_job_token: &apriori_job_token, uuid_idempotency_token: &request.uuid_idempotency_token, maybe_external_third_party_id: &external_job_id, fal_category: FalCategory::ImageGeneration, maybe_inference_args: None, maybe_prompt_token: prompt_token.as_ref(), - maybe_creator_user_token: maybe_user_token, + maybe_creator_user_token: Some(&user_token), maybe_avt_token: maybe_avt_token.as_ref(), creator_ip_address: &ip_address, creator_set_visibility: Visibility::Public,