From 160dbcf91d6c75b8d3a0f6a7e4a5ebe232cdc88e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 1 Mar 2026 04:25:51 -0500 Subject: [PATCH] cost estimate endpoint --- .../cost_estimate/estimate_image_cost.rs | 116 +++++++++++++ .../src/generate/cost_estimate/mod.rs | 1 + .../common/generation/common_image_model.rs | 10 ++ .../public/enums/src/common/generation/mod.rs | 1 + .../image/estimate_image_cost_handler.rs | 157 ++++++++++++++++++ .../generate/cost_estimate/image/mod.rs | 1 + .../endpoints/generate/cost_estimate/mod.rs | 1 + .../application_routes/generate_routes.rs | 5 + 8 files changed, 292 insertions(+) create mode 100644 crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/estimate_image_cost.rs create mode 100644 crates/schema/public/enums/src/common/generation/common_image_model.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/mod.rs diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/estimate_image_cost.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/estimate_image_cost.rs new file mode 100644 index 0000000000..fc5e1ad9c2 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/estimate_image_cost.rs @@ -0,0 +1,116 @@ +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +use enums::common::generation::common_aspect_ratio::CommonAspectRatio; +use enums::common::generation::common_image_model::CommonImageModel; +use enums::common::generation::common_video_resolution::CommonVideoResolution; +use enums::common::generation_provider::GenerationProvider; + +pub const ESTIMATE_IMAGE_COST_PATH: &str = "/v1/generate/cost_estimate/image"; + +/// Request body for the image cost estimate endpoint. +#[derive(Serialize, Deserialize, ToSchema)] +pub struct EstimateImageCostRequest { + /// The image model to estimate costs for. + pub model: CommonImageModel, + + /// The provider to route the generation through. + pub provider: GenerationProvider, + + /// The type of generation (determines whether input images are involved). + /// This is a tagged enum, so it looks like: + /// "generation_mode": {"type": "text_to_image"} + /// "generation_mode": {"type": "image_edit", "count": 2} + pub generation_mode: GenerationMode, + + /// Optional aspect ratio. + pub aspect_ratio: Option, + + /// Optional resolution. + pub resolution: Option, + + /// Number of images to generate in parallel. + pub image_batch_count: Option, +} + +/// Describes the type of image generation being requested. +#[derive(Serialize, Deserialize, ToSchema)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum GenerationMode { + /// No input images — generate from prompt only. + TextToImage, + /// One or more input images are provided for editing/compositing. + ImageEdit { count: u32 }, +} + +/// Response body for the image cost estimate endpoint. +#[derive(Serialize, Deserialize, ToSchema)] +pub struct EstimateImageCostResponse { + pub success: bool, + + /// Estimated cost in credits. + pub cost_in_credits: Option, + + /// Estimated cost in USD cents. + pub cost_in_usd_cents: Option, + + /// Whether the generation is free for this user/plan. + pub is_free: bool, + + /// Whether the user has unlimited generations. + pub is_unlimited: bool, + + /// Whether the user is rate limited. + pub is_rate_limited: bool, + + /// Whether the output will have a watermark. + pub has_watermark: bool, +} + +/// Error response for the image cost estimate endpoint. +#[derive(Serialize, Deserialize, ToSchema)] +pub struct EstimateImageCostError { + pub success: bool, + pub error_type: EstimateImageCostErrorType, + pub error_message: String, +} + +#[derive(Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum EstimateImageCostErrorType { + InvalidProviderForModel, + InvalidInput, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_serialization_text_to_image() { + let request = EstimateImageCostRequest { + model: CommonImageModel::NanaBananaPro, + provider: GenerationProvider::Artcraft, + generation_mode: GenerationMode::TextToImage, + aspect_ratio: None, + resolution: None, + image_batch_count: None, + }; + let serialized = serde_json::to_string(&request).unwrap(); + assert!(serialized.contains("\"generation_mode\":{\"type\":\"text_to_image\"}")); + } + + #[test] + fn test_serialization_image_edit() { + let request = EstimateImageCostRequest { + model: CommonImageModel::NanaBananaPro, + provider: GenerationProvider::Artcraft, + generation_mode: GenerationMode::ImageEdit { count: 2 }, + aspect_ratio: None, + resolution: None, + image_batch_count: None, + }; + let serialized = serde_json::to_string(&request).unwrap(); + assert!(serialized.contains("\"generation_mode\":{\"type\":\"image_edit\",\"count\":2}")); + } +} diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/mod.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/mod.rs index 65a10d3870..e38c09aebf 100644 --- a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/mod.rs +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/cost_estimate/mod.rs @@ -1 +1,2 @@ +pub mod estimate_image_cost; pub mod estimate_video_cost; diff --git a/crates/schema/public/enums/src/common/generation/common_image_model.rs b/crates/schema/public/enums/src/common/generation/common_image_model.rs new file mode 100644 index 0000000000..599c131406 --- /dev/null +++ b/crates/schema/public/enums/src/common/generation/common_image_model.rs @@ -0,0 +1,10 @@ +use utoipa::ToSchema; + +/// Image models available for generation. +/// Mirrors artcraft_router::api::common_image_model::CommonImageModel. +#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum CommonImageModel { + #[serde(rename = "nano_banana_pro")] + NanaBananaPro, +} diff --git a/crates/schema/public/enums/src/common/generation/mod.rs b/crates/schema/public/enums/src/common/generation/mod.rs index 77374dbaa0..260f7ba859 100644 --- a/crates/schema/public/enums/src/common/generation/mod.rs +++ b/crates/schema/public/enums/src/common/generation/mod.rs @@ -1,3 +1,4 @@ pub mod common_aspect_ratio; +pub mod common_image_model; pub mod common_video_model; pub mod common_video_resolution; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs new file mode 100644 index 0000000000..8e9f483a41 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs @@ -0,0 +1,157 @@ +use std::fmt::{Display, Formatter}; + +use actix_http::StatusCode; +use actix_web::web::Json; +use actix_web::{HttpResponse, ResponseError}; +use artcraft_api_defs::generate::cost_estimate::estimate_image_cost::{ + EstimateImageCostError, EstimateImageCostErrorType, EstimateImageCostRequest, + EstimateImageCostResponse, +}; +use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as RouterAspectRatio; +use artcraft_router::api::common_image_model::CommonImageModel as RouterImageModel; +use artcraft_router::api::common_resolution::CommonResolution as RouterResolution; +use artcraft_router::api::provider::Provider as RouterProvider; +use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy; +use artcraft_router::generate::generate_image::generate_image_request::GenerateImageRequest; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio; +use enums::common::generation::common_image_model::CommonImageModel; +use enums::common::generation::common_video_resolution::CommonVideoResolution; +use enums::common::generation_provider::GenerationProvider; + +/// Estimate the credit and USD cost of an image generation request. +/// Does not require authentication and does not charge any credits. +#[utoipa::path( + post, + tag = "Cost Estimate", + path = "/v1/generate/cost_estimate/image", + responses( + (status = 200, description = "Cost estimate", body = EstimateImageCostResponse), + (status = 400, description = "Invalid request", body = EstimateImageCostError), + ), +)] +pub async fn estimate_image_cost_handler( + request: Json, +) -> Result, HandlerError> { + let router_provider = map_provider(request.provider, request.model)?; + let router_model = map_image_model(request.model)?; + let router_aspect_ratio = request.aspect_ratio.map(map_aspect_ratio); + let router_resolution = request.resolution.map(map_resolution); + + let router_request = GenerateImageRequest { + model: router_model, + provider: router_provider, + prompt: None, + image_inputs: None, + resolution: router_resolution, + aspect_ratio: router_aspect_ratio, + image_batch_count: request.image_batch_count, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayLessDowngrade, + idempotency_token: None, + }; + + let plan = router_request.build() + .map_err(|e| HandlerError::InvalidInput(format!("{}", e)))?; + + let estimate = plan.estimate_costs(); + + Ok(Json(EstimateImageCostResponse { + success: true, + cost_in_credits: estimate.cost_in_credits, + cost_in_usd_cents: estimate.cost_in_usd_cents, + is_free: estimate.is_free, + is_unlimited: estimate.is_unlimited, + is_rate_limited: estimate.is_rate_limited, + has_watermark: estimate.has_watermark, + })) +} + +/// Local error type — wraps the serializable API error struct so we can implement +/// ResponseError (orphan rules prevent implementing it directly on the foreign type). +#[derive(Debug)] +pub enum HandlerError { + InvalidProviderForModel { provider: String, model: String }, + InvalidInput(String), +} + +impl std::error::Error for HandlerError {} + +impl Display for HandlerError { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +impl ResponseError for HandlerError { + fn status_code(&self) -> StatusCode { + StatusCode::BAD_REQUEST + } + + fn error_response(&self) -> HttpResponse { + let (error_type, error_message) = match self { + HandlerError::InvalidProviderForModel { provider, model } => ( + EstimateImageCostErrorType::InvalidProviderForModel, + format!("Provider '{}' is not supported for model '{}'", provider, model), + ), + HandlerError::InvalidInput(msg) => ( + EstimateImageCostErrorType::InvalidInput, + msg.clone(), + ), + }; + HttpResponse::BadRequest().json(EstimateImageCostError { + success: false, + error_type, + error_message, + }) + } +} + +fn map_provider( + provider: GenerationProvider, + model: CommonImageModel, +) -> Result { + match provider { + GenerationProvider::Artcraft => Ok(RouterProvider::Artcraft), + other => Err(HandlerError::InvalidProviderForModel { + provider: format!("{:?}", other), + model: format!("{:?}", model), + }), + } +} + +fn map_image_model(model: CommonImageModel) -> Result { + let router_model = match model { + CommonImageModel::NanaBananaPro => RouterImageModel::NanaBananaPro, + }; + Ok(router_model) +} + +fn map_aspect_ratio(ar: CommonAspectRatio) -> RouterAspectRatio { + match ar { + CommonAspectRatio::Auto => RouterAspectRatio::Auto, + CommonAspectRatio::Square => RouterAspectRatio::Square, + CommonAspectRatio::WideThreeByTwo => RouterAspectRatio::WideThreeByTwo, + CommonAspectRatio::WideFourByThree => RouterAspectRatio::WideFourByThree, + CommonAspectRatio::WideFiveByFour => RouterAspectRatio::WideFiveByFour, + CommonAspectRatio::WideSixteenByNine => RouterAspectRatio::WideSixteenByNine, + CommonAspectRatio::WideTwentyOneByNine => RouterAspectRatio::WideTwentyOneByNine, + CommonAspectRatio::TallTwoByThree => RouterAspectRatio::TallTwoByThree, + CommonAspectRatio::TallThreeByFour => RouterAspectRatio::TallThreeByFour, + CommonAspectRatio::TallFourByFive => RouterAspectRatio::TallFourByFive, + CommonAspectRatio::TallNineBySixteen => RouterAspectRatio::TallNineBySixteen, + CommonAspectRatio::TallNineByTwentyOne => RouterAspectRatio::TallNineByTwentyOne, + CommonAspectRatio::Wide => RouterAspectRatio::Wide, + CommonAspectRatio::Tall => RouterAspectRatio::Tall, + CommonAspectRatio::Auto2k => RouterAspectRatio::Auto2k, + CommonAspectRatio::Auto4k => RouterAspectRatio::Auto4k, + CommonAspectRatio::SquareHd => RouterAspectRatio::SquareHd, + } +} + +fn map_resolution(res: CommonVideoResolution) -> RouterResolution { + match res { + CommonVideoResolution::OneK => RouterResolution::OneK, + CommonVideoResolution::TwoK => RouterResolution::TwoK, + CommonVideoResolution::ThreeK => RouterResolution::ThreeK, + CommonVideoResolution::FourK => RouterResolution::FourK, + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/mod.rs new file mode 100644 index 0000000000..44b14d777c --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/mod.rs @@ -0,0 +1 @@ +pub mod estimate_image_cost_handler; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/mod.rs index b04b953edb..c10a6ca15e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/mod.rs @@ -1 +1,2 @@ +pub mod image; pub mod video; diff --git a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs index de66ae4996..2eb4bcb6c5 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs @@ -1,3 +1,4 @@ +use crate::http_server::endpoints::generate::cost_estimate::image::estimate_image_cost_handler::estimate_image_cost_handler; use crate::http_server::endpoints::generate::cost_estimate::video::estimate_video_cost_handler::estimate_video_cost_handler; use crate::http_server::endpoints::generate::image::edit::flux_pro_kontext_max_edit_image_handler::flux_pro_kontext_max_edit_image_handler; use crate::http_server::endpoints::generate::image::edit::gemini_25_flash_edit_image_handler::gemini_25_flash_edit_image_handler; @@ -53,6 +54,10 @@ where { app.service(web::scope("/v1/generate") .service(web::scope("/cost_estimate") + .service(web::resource("/image") + .route(web::post().to(estimate_image_cost_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) .service(web::resource("/video") .route(web::post().to(estimate_video_cost_handler)) .route(web::head().to(|| HttpResponse::Ok()))