Compare commits

...
13 changed files with 12 additions and 798 deletions
Generated
+10 -12
View File
@@ -2439,7 +2439,7 @@ checksum = "2eac901828f88a5241ee0600950ab981148a18f2f756900ffba1b125ca6a3ef9"
dependencies = [
"cookie 0.18.1",
"document-features",
"idna 1.1.0",
"idna 1.0.3",
"log",
"publicsuffix",
"serde",
@@ -4109,9 +4109,9 @@ checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b"
[[package]]
name = "form_urlencoded"
version = "1.2.2"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf"
checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456"
dependencies = [
"percent-encoding",
]
@@ -5404,9 +5404,9 @@ dependencies = [
[[package]]
name = "idna"
version = "1.1.0"
version = "1.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e"
dependencies = [
"idna_adapter",
"smallvec 1.15.1",
@@ -8301,7 +8301,7 @@ version = "2.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f42ea446cab60335f76979ec15e12619a2165b5ae2c12166bef27d283a9fadf"
dependencies = [
"idna 1.1.0",
"idna 1.0.3",
"psl-types",
]
@@ -10599,7 +10599,6 @@ dependencies = [
"mimetypes",
"mysql_queries",
"once_cell",
"openai_sora_client",
"password",
"primitives",
"prompt_engineering",
@@ -10616,7 +10615,6 @@ dependencies = [
"serde_derive",
"serde_json",
"server_environment",
"shared_service_components",
"speculoos",
"sqlx",
"strum",
@@ -12435,12 +12433,12 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
[[package]]
name = "url"
version = "2.5.7"
version = "2.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08bc136a29a3d1758e07a9cca267be308aeebf5cfd5a10f3f67ab2097683ef5b"
checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60"
dependencies = [
"form_urlencoded",
"idna 1.1.0",
"idna 1.0.3",
"percent-encoding",
"serde",
]
@@ -13731,7 +13729,7 @@ dependencies = [
"hyper-rustls",
"hyper-tls 0.6.0",
"hyper-util",
"idna 1.1.0",
"idna 1.0.3",
"indexmap 1.9.3",
"indexmap 2.9.0",
"lazy_static",
@@ -64,7 +64,7 @@ memory_caching = { path = "../../../lib/caching/memory_caching" }
migration = { path = "../../../lib/migration" }
mimetypes.workspace = true
mysql_queries = { path = "../../../schema/database/mysql_queries" }
openai_sora_client = { path = "../../../api_clients/openai_sora_client" }
#openai_sora_client = { path = "../../../api_clients/openai_sora_client" }
password = { path = "../../../lib/password" }
primitives = { path = "../../../lib/data/primitives" }
prompt_engineering = { path = "../../../lib/machine_learning/prompt_engineering" }
@@ -73,7 +73,7 @@ redis_common = { path = "../../../schema/database/redis_common" }
redis_schema = { path = "../../../schema/database/redis_schema" }
reusable_types = { path = "../../../schema/public/reusable_types" }
server_environment = { path = "../../../lib/server_environment" }
shared_service_components.workspace = true
#shared_service_components.workspace = true
thumbnail_generator = { path = "../../../api_clients/thumbnail_generator" }
tokens.workspace = true
tts_common = { path = "../../../lib/tts_common" }
@@ -1 +0,0 @@
pub mod set_sora_secret_handler;
@@ -1,137 +0,0 @@
use std::fmt;
use std::sync::Arc;
use actix_web::error::ResponseError;
use actix_web::http::StatusCode;
use actix_web::web::{Form, Json};
use actix_web::{web, HttpRequest, HttpResponse};
use log::error;
use r2d2_redis::redis::Commands;
use utoipa::ToSchema;
use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
use shared_service_components::sora_redis_credentials::get_sora_credentials_from_redis::get_sora_credentials_from_redis;
use shared_service_components::sora_redis_credentials::keys::RedisSoraCredentialSubkey;
use shared_service_components::sora_redis_credentials::set_sora_credential_field_in_redis::set_sora_credential_field_in_redis;
use crate::state::server_state::ServerState;
#[derive(Deserialize, Clone, Copy)]
#[serde(rename_all = "snake_case")]
pub enum SetSoraSecretType {
Bearer,
Cookie,
Sentinel
}
#[derive(Deserialize, ToSchema)]
pub struct SetSoraSecretRequest {
pub key: SetSoraSecretType,
pub value: String,
}
#[derive(Serialize, ToSchema)]
pub struct SetSoraSecretSuccessResponse {
pub success: bool,
pub bearer: String,
pub cookie: String,
pub sentinel: String,
}
// =============== Error Response ===============
#[derive(Debug, Serialize, ToSchema)]
pub enum SetSoraSecretError {
BadInput(String),
NotAuthorized,
ServerError,
}
impl ResponseError for SetSoraSecretError {
fn status_code(&self) -> StatusCode {
match *self {
SetSoraSecretError::BadInput(_) => StatusCode::BAD_REQUEST,
SetSoraSecretError::NotAuthorized => StatusCode::UNAUTHORIZED,
SetSoraSecretError::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 SetSoraSecretError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self)
}
}
// =============== Handler ===============
/// Set a Sora secret
#[utoipa::path(
post,
tag = "Control Plane",
path = "/v1/control_plane/sora_secret",
responses(
(status = 200, description = "Success", body = SetSoraSecretSuccessResponse),
(status = 400, description = "Bad input", body = SetSoraSecretError),
(status = 401, description = "Not authorized", body = SetSoraSecretError),
(status = 500, description = "Server error", body = SetSoraSecretError),
),
params(
("request" = SetSoraSecretRequest, description = "Payload for Request"),
("path" = MediaFileTokenPathInfo, description = "Path for Request")
)
)]
pub async fn set_sora_secret_handler(
http_request: HttpRequest,
request: Form<SetSoraSecretRequest>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<SetSoraSecretSuccessResponse>, SetSoraSecretError> {
// TODO(bt,2025-04-03): Some form of authentication should be required for this endpoint.
let value = request.value.trim().to_string();
let mut redis = server_state.redis_pool
.get()
.map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
set_sora_credential_field_in_redis(
&mut redis,
match request.key {
SetSoraSecretType::Bearer => RedisSoraCredentialSubkey::Bearer,
SetSoraSecretType::Cookie => RedisSoraCredentialSubkey::Cookie,
SetSoraSecretType::Sentinel => RedisSoraCredentialSubkey::Sentinel,
},
&value
).map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
let credentials = get_sora_credentials_from_redis(&mut redis)
.map_err(|e| {
error!("redis error: {:?}", e);
SetSoraSecretError::ServerError
})?;
// TODO(bt,2025-04-03): Probably not smart to return, but useful for debugging.
Ok(Json(SetSoraSecretSuccessResponse {
success: true,
cookie: credentials.cookies.as_str().to_string(),
bearer: credentials.jwt_bearer_token
.as_ref()
.map(|b| b.as_str().to_string())
.unwrap_or_else(|| "".to_string()),
sentinel: credentials.sora_sentinel
.as_ref()
.map(|s| s.as_str().to_string())
.unwrap_or_else(|| "".to_string()),
}))
}
@@ -1,7 +1,6 @@
pub mod animation;
pub mod api_tokens;
pub mod categories;
pub mod control_plane;
pub mod conversion;
pub mod engine;
pub mod events;
@@ -1,69 +0,0 @@
use actix_web::HttpRequest;
use anyhow::anyhow;
use errors::AnyhowResult;
use http_server_common::request::get_request_header_optional::get_request_header_optional;
use log::{error, info};
use openai_sora_client::creds::sora_credential_builder::SoraCredentialBuilder;
use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
use r2d2_redis::r2d2::PooledConnection;
use r2d2_redis::RedisConnectionManager;
use shared_service_components::sora_redis_credentials::get_sora_credentials_from_redis::get_sora_credentials_from_redis;
const SORA_BEARER_HEADER_NAME : &str = "sora-bearer";
const SORA_COOKIE_HEADER_NAME : &str = "sora-cookie";
const SORA_SENTINEL_HEADER_NAME : &str = "sora-sentinel";
pub fn get_sora_credentials_from_request(
http_request: &HttpRequest,
redis: &mut PooledConnection<RedisConnectionManager>,
) -> AnyhowResult<SoraCredentialSet> {
let sora_bearer = get_request_header_optional(&http_request, SORA_BEARER_HEADER_NAME);
let sora_cookie = get_request_header_optional(&http_request, SORA_COOKIE_HEADER_NAME);
let sora_sentinel = get_request_header_optional(&http_request, SORA_SENTINEL_HEADER_NAME);
match (&sora_bearer, &sora_cookie, &sora_sentinel) {
(Some(bearer), Some(cookie), Some(sentinel)) => {
info!("Sora credentials were present in HTTP request headers!");
let mut creds = SoraCredentialBuilder::new()
.with_cookies(cookie)
.with_jwt_bearer_token(bearer)
.with_sora_sentinel(sentinel)
.build()?;
return Ok(creds);
}
_ => {} // Fall through
}
let redis_result = get_sora_credentials_from_redis(redis);
match redis_result {
Ok(credentials) => {
let mut builder = credentials.to_builder();
if let Some(bearer) = sora_bearer {
info!("override bearer value");
builder = builder.with_jwt_bearer_token(&bearer);
}
if let Some(cookie) = sora_cookie {
info!("override cookie value");
builder = builder.with_cookies(&cookie)
}
if let Some(sentinel) = sora_sentinel {
info!("override sentinel value");
builder = builder.with_sora_sentinel(&sentinel);
}
return Ok(builder.build()?);
},
Err(err) => {
error!("Failed to get Sora credentials from Redis; using env as backup. Error: {:?}", err);
// NB: Fall through
}
};
// TODO(bt): I think I can kill this entire module. We don't call the API from our backend.
error!("Failed to load Sora credentials from environment");
Err(anyhow!("Failed to load Sora credentials from environment"))
}
@@ -1,4 +1,2 @@
pub mod upload;
pub mod prompt;
pub mod get_sora_credentials_from_request;
pub mod update_gpt_image_job_status_handler;
@@ -1,541 +0,0 @@
#![forbid(unused_mut)]
use std::fmt::Debug;
use std::fmt::{Display, Formatter};
use std::path::PathBuf;
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 bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use cloud_storage::bucket_client::BucketClient;
use log::warn;
use log::{debug, error};
use mysql_queries::queries::media_files::get::get_media_file::get_media_file;
use openai_sora_client::requests::auth_sentinel::generate_sentinel_token::generate_sentinel_token;
use serde::Deserialize;
use serde::Serialize;
use shared_service_components::sora_redis_credentials::keys::RedisSoraCredentialSubkey;
use shared_service_components::sora_redis_credentials::set_sora_credential_field_in_redis::set_sora_credential_field_in_redis;
use shared_service_components::sora_redis_sentinel_refresh::refresh::refresh_sentinel;
use sqlx::MySqlPool;
use tempdir::TempDir;
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 enums::common::visibility::Visibility;
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::sora_image_gen_args::SoraImageGenArgs;
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 openai_sora_client::error::sora_error::SoraError;
use openai_sora_client::error::sora_specific_api_error::SoraSpecificApiError;
use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
use openai_sora_client::requests::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest};
use openai_sora_client::requests::upload::upload_media_from_file::sora_media_upload_from_file;
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::users::UserToken;
use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session;
use crate::http_server::endpoints::image_studio::get_sora_credentials_from_request::get_sora_credentials_from_request;
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 = 1;
/// The maximum number of images (batch size) to generate for each request.
const MAXIMUM_IMAGE_COUNT: u32 = 2;
/// 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(Deserialize, ToSchema)]
pub struct EnqueueStudioImageGenRequest {
pub uuid_idempotency_token: String,
/// Image media file; the engine or canvas snapshot (screenshot).
pub snapshot_media_token: MediaFileToken,
/// The user's image generation prompt.
pub prompt: String,
/// Turn off the system prompt.
pub disable_system_prompt: Option<bool>,
/// Additional images to include (optional). Up to nine images.
pub additional_images: Option<Vec<MediaFileToken>>,
#[deprecated(note="use `additional_images` instead. if both are present, `additional_images` will be used.")]
pub maybe_additional_images: Option<Vec<MediaFileToken>>,
pub maybe_number_of_samples: Option<u32>,
}
#[derive(Serialize, ToSchema)]
pub struct EnqueueImageGenRequestSuccessResponse {
pub success: bool,
pub inference_job_token: InferenceJobToken,
}
#[derive(Debug, ToSchema)]
pub enum EnqueueImageGenRequestError {
BadInput(String),
NotAuthorized,
ServerError,
RateLimited,
TooManyConcurrentTasks,
}
impl ResponseError for EnqueueImageGenRequestError {
fn status_code(&self) -> StatusCode {
match *self {
EnqueueImageGenRequestError::BadInput(_) => StatusCode::BAD_REQUEST,
EnqueueImageGenRequestError::NotAuthorized => StatusCode::UNAUTHORIZED,
EnqueueImageGenRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
EnqueueImageGenRequestError::RateLimited => StatusCode::TOO_MANY_REQUESTS,
EnqueueImageGenRequestError::TooManyConcurrentTasks => StatusCode::TOO_MANY_REQUESTS,
}
}
fn error_response(&self) -> HttpResponse {
let error_reason = match self {
EnqueueImageGenRequestError::BadInput(reason) => reason.to_string(),
EnqueueImageGenRequestError::NotAuthorized => "unauthorized".to_string(),
EnqueueImageGenRequestError::ServerError => "server error".to_string(),
EnqueueImageGenRequestError::RateLimited => "rate limited".to_string(),
EnqueueImageGenRequestError::TooManyConcurrentTasks => "too many concurrent tasks".to_string(),
};
to_simple_json_error(&error_reason, self.status_code())
}
}
impl Display for EnqueueImageGenRequestError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?}", self)
}
}
/// Prompt image generation using image studio
#[utoipa::path(
post,
tag = "Image Studio",
path = "/v1/image_studio/prompt",
params(
("request" = EnqueueStudioImageGenRequest, description = "Payload for Image Generation Request"),
),
responses(
(status = 200, description = "Success response", body = EnqueueImageGenRequestSuccessResponse),
(status = 400, description = "Bad input", body = EnqueueImageGenRequestError),
(status = 401, description = "Not authorized", body = EnqueueImageGenRequestError),
(status = 429, description = "Rate limited", body = EnqueueImageGenRequestError),
(status = 500, description = "Server error", body = EnqueueImageGenRequestError),
),
)]
pub async fn enqueue_studio_image_generation_handler(http_request: HttpRequest, request: Json<EnqueueStudioImageGenRequest>, server_state: Data<Arc<ServerState>>) -> Result<HttpResponse, EnqueueImageGenRequestError> {
validate_request(&request)?;
let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| {
warn!("MySql pool error: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
// ==================== 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);
EnqueueImageGenRequestError::ServerError
})?;
let mut maybe_user_token: Option<UserToken> = None;
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_old, 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(EnqueueImageGenRequestError::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) {
return Err(EnqueueImageGenRequestError::RateLimited);
}
// Get up IP address
let ip_address = get_request_ip(&http_request);
// ==================== HANDLE IDEMPOTENCY ==================== //
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(EnqueueImageGenRequestError::BadInput(reason));
}
insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection).await.map_err(|err| {
error!("Error inserting idempotency token: {:?}", err);
EnqueueImageGenRequestError::BadInput("invalid idempotency token".to_string())
})?;
// ==================== DOWNLOAD FILES AND UPLOAD TO SORA ==================== //
// TODO: Maybe this moves to a job.
let work_temp_dir = server_state.temp_dir_creator.new_tempdir("image_studio").unwrap();
let public_bucket_client = server_state.public_bucket_client.clone();
let scene_media_token = request.snapshot_media_token.clone();
let scene_media_path = query_and_download_media_file(&scene_media_token, &work_temp_dir, &public_bucket_client, &server_state.mysql_pool).await.map_err(|err| {
error!("Failed to download scene media file: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
let additional_images;
// NB: Frontend is calling us with the wrong field name. Let's just accept both.
if let Some(images) = &request.additional_images {
additional_images = images.clone();
} else if let Some(images) = &request.maybe_additional_images {
additional_images = images.clone();
} else {
additional_images = vec![];
}
let mut files_to_upload = Vec::with_capacity(1 + additional_images.len());
files_to_upload.push(scene_media_path.clone());
debug!("Additional images to download: {}", additional_images.len());
for (i, media_file_token ) in additional_images.iter().enumerate() {
debug!("Downloading additional image {} of {} ...", (i+1), additional_images.len());
let media_file_path = query_and_download_media_file(&media_file_token, &work_temp_dir, &public_bucket_client, &server_state.mysql_pool).await.map_err(|err| {
error!("Failed to download additional media file: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
files_to_upload.push(media_file_path.clone());
}
// ==================== HANDLE SORA CREDENTIALS ==================== //
let mut redis = server_state.redis_pool.get().map_err(|e| {
error!("redis error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
let sora_credentials = get_sora_credentials_from_request(&http_request, &mut redis).map_err(|e| {
error!("sora credential error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
// ==================== HANDLE SORA UPLOAD ==================== //
let mut sora_media_tokens = Vec::with_capacity(files_to_upload.len());
for (i, file_path) in files_to_upload.iter().enumerate() {
debug!("Uploading file {} of {} to Sora...", (i+1), files_to_upload.len());
let sora_upload_response =
sora_media_upload_from_file(file_path, &sora_credentials, None)
.await
.map_err(|err| {
error!("Failed to upload scene media to Sora: {:?}", err);
EnqueueImageGenRequestError::ServerError
})?;
debug!("Uploaded media to Sora : {:?}", sora_upload_response);
sora_media_tokens.push(sora_upload_response.id);
}
// ==================== HANDLE SORA PROMPT ==================== //
let number_of_samples = match request.maybe_number_of_samples {
Some(val) => {
if val > MAXIMUM_IMAGE_COUNT {
MAXIMUM_IMAGE_COUNT
} else if val < MINIMUM_IMAGE_COUNT {
MINIMUM_IMAGE_COUNT
} else {
val
}
},
None => 1,
};
let prompt = create_prompt(&request);
debug!("Sending Sora remix request with {} media tokens...", sora_media_tokens.len());
let mut response = sora_image_gen_remix(SoraImageGenRemixRequest {
prompt: prompt.clone(),
num_images: NumImages::One,
image_size: ImageSize::Square,
sora_media_tokens: sora_media_tokens.clone(),
credentials: &sora_credentials,
request_timeout: None,
}).await;
debug!("Sora image gen response: {:?}", response);
if let Err(SoraError::ApiSpecific(SoraSpecificApiError::SentinelBlockError)) = &response {
error!("Sora sentinel block, attempting refresh");
match generate_sentinel_token().await {
Ok(sentinel_token) => {
set_sora_credential_field_in_redis(
&mut redis,
RedisSoraCredentialSubkey::Sentinel,
&sentinel_token
).map_err(|e| {
EnqueueImageGenRequestError::ServerError
})?;
let updated_sora_credentials = get_sora_credentials_from_request(&http_request, &mut redis).map_err(|e| {
error!("sora credential error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
// Retry the request with new sentinel
response = sora_image_gen_remix(SoraImageGenRemixRequest {
prompt: prompt.clone(),
num_images: NumImages::One,
image_size: ImageSize::Square,
sora_media_tokens: sora_media_tokens.clone(),
credentials: &updated_sora_credentials,
request_timeout: None,
}).await;
},
Err(e) => {
// Fallback to the old refresh stragegy relying on the external service
match refresh_sentinel().await {
Ok(new_sentinel) => {
// Update the credentials with new sentinel
let updated_sora_credentials = get_sora_credentials_from_request(&http_request, &mut redis).map_err(|e| {
error!("sora credential error: {:?}", e);
EnqueueImageGenRequestError::ServerError
})?;
// Retry the request with new sentinel
response = sora_image_gen_remix(SoraImageGenRemixRequest {
prompt: prompt.clone(),
num_images: NumImages::One,
image_size: ImageSize::Square,
sora_media_tokens: sora_media_tokens.clone(),
credentials: &updated_sora_credentials,
request_timeout: None,
}).await;
},
Err(e) => {
error!("Failed to refresh Sora sentinel: {:?}", e);
}
}
}
}
}
let response = response.map_err(|err| {
match err {
SoraError::ApiSpecific(SoraSpecificApiError::TokenExpiredError) => {
error!("Sora token expired, needs refresh");
// TODO: Implement token refresh logic here
EnqueueImageGenRequestError::ServerError
},
SoraError::ApiSpecific(SoraSpecificApiError::SentinelBlockError) => {
error!("Sora sentinel block still occurring after refresh");
EnqueueImageGenRequestError::ServerError
},
SoraError::ApiSpecific(SoraSpecificApiError::TooManyConcurrentTasks) => {
error!("Sora too many concurrent tasks");
EnqueueImageGenRequestError::TooManyConcurrentTasks
}
_ => {
error!("Failed to call Sora image generation: {:?}", err);
EnqueueImageGenRequestError::ServerError
},
}
})?;
debug!("Sora image gen response: {:?}", response);
// Store the actual inference args that will go to the database
let inference_args = SoraImageGenArgs {
prompt: Some(prompt),
scene_snapshot_media_token: Some(request.snapshot_media_token.clone()),
maybe_additional_media_file_tokens: Some(additional_images),
maybe_number_of_samples: Some(number_of_samples),
maybe_sora_media_upload_tokens: Some(sora_media_tokens),
maybe_sora_task_id: Some(response.task_id.0)
};
// 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::ImageGenApi,
maybe_product_category: None, // This is not a product anymore
inference_category: InferenceCategory::ImageGeneration,
maybe_model_type: Some(InferenceModelType::ImageGenApi), // 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: None,
maybe_cover_image_media_file_token: None,
maybe_raw_inference_text: None,
maybe_max_duration_seconds: None,
maybe_inference_args: Some(GenericInferenceArgs {
inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration),
args: Some(PolymorphicInferenceArgs::Sg(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: Visibility::Public,
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(EnqueueImageGenRequestError::BadInput("Duplicate idempotency token".to_string()));
}
return Err(EnqueueImageGenRequestError::ServerError);
},
};
let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse { success: true, inference_job_token: job_token };
let body = serde_json::to_string(&response).map_err(|_e| EnqueueImageGenRequestError::ServerError)?;
// Error handling 101 rust result type returned like so.
Ok(HttpResponse::Ok().content_type("application/json").body(body))
}
fn validate_request(request: &Json<EnqueueStudioImageGenRequest>) -> Result<(), EnqueueImageGenRequestError> {
Ok(())
}
#[derive(Debug)]
pub enum DownloadMediaFileError {
IoError(std::io::Error),
MediaFileNotFound,
Other(anyhow::Error),
}
impl From<anyhow::Error> for DownloadMediaFileError {
fn from(err: anyhow::Error) -> Self {
DownloadMediaFileError::Other(err)
}
}
async fn query_and_download_media_file(media_file_token: &MediaFileToken, download_dir: &TempDir, bucket_client: &BucketClient, mysql_pool: &MySqlPool) -> Result<PathBuf, DownloadMediaFileError> {
let media_file_result = get_media_file(&media_file_token, false, mysql_pool).await;
let media_file = match media_file_result {
Ok(Some(result)) => result,
Ok(None) => {
warn!("could not find media file in database: {:?}", media_file_token);
return Err(DownloadMediaFileError::MediaFileNotFound);
},
Err(e) => {
error!("could not query media file: {:?}", e);
return Err(DownloadMediaFileError::Other(e));
},
};
let media_file_bucket_path = MediaFileBucketPath::from_object_hash(
&media_file.public_bucket_directory_hash,
media_file.maybe_public_bucket_prefix.as_deref(),
media_file.maybe_public_bucket_extension.as_deref());
let extension = media_file.maybe_public_bucket_extension.clone().unwrap_or_else(|| ".jpg".to_string());
let download_filename = format!("download_file_{}.{}", media_file.token.as_str(), extension);
let download_file_path = download_dir.path().join(download_filename);
debug!("Downloading from bucket...");
bucket_client.download_file_to_disk(&media_file_bucket_path.to_full_object_pathbuf(), &download_file_path).await?;
Ok(download_file_path)
}
const SYSTEM_PROMPT_SINGLE_IMAGE: &str =
"Use the image as a reference for the composition exactly. Copy the character pose, layout, scene, and the objects exactly - no deviation.";
const SYSTEM_PROMPT_MULTI_IMAGE: &str =
"Use the first image as a reference for the composition exactly. Copy the character pose, layout, scene, and the objects exactly from the first image - no deviation.";
fn create_prompt(request: &EnqueueStudioImageGenRequest) -> String {
let disable_system_prompt = request.disable_system_prompt
.clone()
.unwrap_or(false);
if disable_system_prompt {
return request.prompt.clone();
}
let extra_image_count = request.maybe_additional_images
.as_ref()
.map(|v| v.len())
.unwrap_or(0);
if extra_image_count == 0 {
format!("{} {}", SYSTEM_PROMPT_SINGLE_IMAGE, request.prompt)
} else {
format!("{} {}", SYSTEM_PROMPT_MULTI_IMAGE, request.prompt)
}
}
@@ -1 +0,0 @@
pub mod enqueue_studio_image_generation_handler;
@@ -68,7 +68,6 @@ use crate::http_server::endpoints::voice_designer::voices::list_voices_by_user::
use crate::http_server::endpoints::voice_designer::voices::search_voices::search_voices;
use crate::http_server::endpoints::voice_designer::voices::update_voice::update_voice_handler;
use crate::http_server::routes::legacy_routes::beta_key_routes::add_beta_key_routes;
use crate::http_server::routes::legacy_routes::control_plane_sora_routes::add_control_plane_sora_routes;
use crate::http_server::routes::legacy_routes::desktop_vc_app_routes::add_desktop_vc_app_routes;
use crate::http_server::routes::legacy_routes::image_studio_routes::add_image_studio_routes;
use crate::http_server::routes::legacy_routes::model_download_routes::add_model_download_routes;
@@ -109,7 +108,6 @@ where
app = add_model_download_routes(app);
app = add_studio_gen2_routes(app);
app = add_image_studio_routes(app);
app = add_control_plane_sora_routes(app); // /v1/control_plane/...
app = add_workflow_routes(app);
app = add_engine_routes(app); // /v1/engine/...
app = add_desktop_vc_app_routes(app); // /v1/vc/...
@@ -1,24 +0,0 @@
use crate::http_server::deprecated_endpoints::control_plane::set_sora_secret_handler::set_sora_secret_handler;
use actix_http::body::MessageBody;
use actix_service::ServiceFactory;
use actix_web::dev::{ServiceRequest, ServiceResponse};
use actix_web::{web, App, Error, HttpResponse};
pub fn add_control_plane_sora_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/control_plane")
.service(web::resource("/sora_secret")
.route(web::post().to(set_sora_secret_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
)
}
@@ -1,4 +1,3 @@
use crate::http_server::endpoints::image_studio::prompt::enqueue_studio_image_generation_handler::enqueue_studio_image_generation_handler;
use crate::http_server::endpoints::image_studio::update_gpt_image_job_status_handler::update_gpt_image_job_status_handler;
use crate::http_server::endpoints::image_studio::upload::upload_snapshot_media_file_handler::upload_snapshot_media_file_handler;
use actix_http::body::MessageBody;
@@ -22,10 +21,6 @@ where
.route(web::post().to(upload_snapshot_media_file_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/prompt")
.route(web::post().to(enqueue_studio_image_generation_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/update_job_status")
.route(web::post().to(update_gpt_image_job_status_handler))
.route(web::head().to(|| HttpResponse::Ok()))
@@ -1,6 +1,5 @@
pub (super) mod add_legacy_routes;
mod beta_key_routes;
mod control_plane_sora_routes;
mod image_studio_routes;
mod model_download_routes;
mod studio_gen2_routes;