estimate costs of splats

This commit is contained in:
Brandon Thomas committed 2026-03-09 14:45:45 -04:00
1 parent e070028bf1
commit 7331fa2fcf
8 files changed
+205

No files matched your search

@@ -0,0 +1,59 @@
use serde::{Deserialize, Serialize};
use utoipa::ToSchema;
use enums::common::generation::common_splat_model::CommonSplatModel;
use enums::common::generation_provider::GenerationProvider;
pub const ESTIMATE_SPLAT_COST_PATH: &str = "/v1/generate/cost_estimate/splat";
/// Request body for the splat cost estimate endpoint.
#[derive(Serialize, Deserialize, ToSchema)]
pub struct EstimateSplatCostRequest {
/// The splat model to estimate costs for.
pub model: CommonSplatModel,
/// The provider to route the generation through.
pub provider: GenerationProvider,
/// Whether the request includes a reference image.
pub has_reference_image: Option<bool>,
}
/// Response body for the splat cost estimate endpoint.
#[derive(Serialize, Deserialize, ToSchema)]
pub struct EstimateSplatCostResponse {
pub success: bool,
/// Estimated cost in credits.
pub cost_in_credits: Option<u64>,
/// Estimated cost in USD cents.
pub cost_in_usd_cents: Option<u64>,
/// 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 splat cost estimate endpoint.
#[derive(Serialize, Deserialize, ToSchema)]
pub struct EstimateSplatCostError {
pub success: bool,
pub error_type: EstimateSplatCostErrorType,
pub error_message: String,
}
#[derive(Serialize, Deserialize, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum EstimateSplatCostErrorType {
InvalidProviderForModel,
InvalidInput,
}
@@ -1,2 +1,3 @@
pub mod estimate_image_cost;
pub mod estimate_splat_cost;
pub mod estimate_video_cost;
@@ -0,0 +1,13 @@
use utoipa::ToSchema;
/// Splat models available for generation.
/// Mirrors artcraft_router::api::common_splat_model::CommonSplatModel.
#[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)]
#[serde(rename_all = "snake_case")]
pub enum CommonSplatModel {
#[serde(rename = "marble_0p1_mini")]
Marble0p1Mini,
#[serde(rename = "marble_0p1_plus")]
Marble0p1Plus,
}
@@ -1,4 +1,5 @@
pub mod common_aspect_ratio;
pub mod common_image_model;
pub mod common_splat_model;
pub mod common_video_model;
pub mod common_video_resolution;
@@ -1,2 +1,3 @@
pub mod image;
pub mod splat;
pub mod video;
@@ -0,0 +1,124 @@
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_splat_cost::{
EstimateSplatCostError, EstimateSplatCostErrorType, EstimateSplatCostRequest,
EstimateSplatCostResponse,
};
use artcraft_router::api::common_splat_model::CommonSplatModel as RouterSplatModel;
use artcraft_router::api::provider::Provider as RouterProvider;
use artcraft_router::generate::generate_splat::generate_splat_request::GenerateSplatRequest;
use enums::common::generation::common_splat_model::CommonSplatModel;
use enums::common::generation_provider::GenerationProvider;
/// Estimate the credit and USD cost of a splat generation request.
/// Does not require authentication and does not charge any credits.
#[utoipa::path(
post,
tag = "Cost Estimate",
path = "/v1/generate/cost_estimate/splat",
responses(
(status = 200, description = "Cost estimate", body = EstimateSplatCostResponse),
(status = 400, description = "Invalid request", body = EstimateSplatCostError),
),
)]
pub async fn estimate_splat_cost_handler(
request: Json<EstimateSplatCostRequest>,
) -> Result<Json<EstimateSplatCostResponse>, HandlerError> {
let router_provider = map_provider(request.provider, request.model)?;
let router_model = map_splat_model(request.model);
let router_request = GenerateSplatRequest {
model: router_model,
provider: router_provider,
prompt: None,
reference_images: if request.has_reference_image.unwrap_or(false) {
// NB: The cost estimation only cares whether a reference image is present,
// not the actual image data. We pass None here; the plan will detect
// "has reference image" from the model's cost function.
// However, the build() path requires no images to estimate cost.
None
} else {
None
},
idempotency_token: None,
};
let plan = router_request.build()
.map_err(|e| HandlerError::InvalidInput(format!("{}", e)))?;
let estimate = plan.estimate_costs();
Ok(Json(EstimateSplatCostResponse {
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 } => (
EstimateSplatCostErrorType::InvalidProviderForModel,
format!("Provider '{}' is not supported for model '{}'", provider, model),
),
HandlerError::InvalidInput(msg) => (
EstimateSplatCostErrorType::InvalidInput,
msg.clone(),
),
};
HttpResponse::BadRequest().json(EstimateSplatCostError {
success: false,
error_type,
error_message,
})
}
}
fn map_provider(
provider: GenerationProvider,
model: CommonSplatModel,
) -> Result<RouterProvider, HandlerError> {
match provider {
GenerationProvider::Artcraft => Ok(RouterProvider::Artcraft),
other => Err(HandlerError::InvalidProviderForModel {
provider: format!("{:?}", other),
model: format!("{:?}", model),
}),
}
}
fn map_splat_model(model: CommonSplatModel) -> RouterSplatModel {
match model {
CommonSplatModel::Marble0p1Mini => RouterSplatModel::Marble0p1Mini,
CommonSplatModel::Marble0p1Plus => RouterSplatModel::Marble0p1Plus,
}
}
@@ -0,0 +1 @@
pub mod estimate_splat_cost_handler;
@@ -1,4 +1,5 @@
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::splat::estimate_splat_cost_handler::estimate_splat_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;
@@ -65,6 +66,10 @@ where
.route(web::post().to(estimate_image_cost_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/splat")
.route(web::post().to(estimate_splat_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()))