From ff360f83b133c9493d3fcab2de7160bbd2ce7ec0 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 24 Jul 2025 00:21:00 -0400
Subject: [PATCH] populate provider and model

---
 .../src/queries/prompts/insert_prompt.rs             | 12 ++++++++++++
 .../image_generation/sd/process_job_inference.rs     |  4 ++++
 .../steps/validate_and_save_results.rs               |  4 ++++
 .../image/edit/gpt_image_1_edit_image_handler.rs     |  4 ++++
 .../video/generate_veo_2_image_to_video_handler.rs   |  4 ++++
 5 files changed, 28 insertions(+)

diff --git a/crates/schema/database/mysql_queries/src/queries/prompts/insert_prompt.rs b/crates/schema/database/mysql_queries/src/queries/prompts/insert_prompt.rs
index 110f1dde55..06471a36fa 100644
--- a/crates/schema/database/mysql_queries/src/queries/prompts/insert_prompt.rs
+++ b/crates/schema/database/mysql_queries/src/queries/prompts/insert_prompt.rs
@@ -6,6 +6,8 @@ use sqlx;
 use sqlx::{Executor, MySql};
 
 use enums::by_table::prompts::prompt_type::PromptType;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
 use errors::AnyhowResult;
 use tokens::tokens::prompts::PromptToken;
 use tokens::tokens::users::UserToken;
@@ -21,6 +23,10 @@ pub struct InsertPromptArgs<'e, 'c,  E>
   pub prompt_type: PromptType,
 
   pub maybe_creator_user_token: Option<&'e UserToken>,
+  
+  pub maybe_model_type: Option<ModelType>,
+  
+  pub maybe_generation_provider: Option<GenerationProvider>,
 
   pub maybe_positive_prompt: Option<&'e str>,
 
@@ -65,9 +71,13 @@ SET
   prompt_type = ?,
 
   maybe_creator_user_token = ?,
+  
+  maybe_model_type = ?,
+  maybe_generation_provider = ?,
 
   maybe_positive_prompt = ?,
   maybe_negative_prompt = ?,
+  
   maybe_other_args = ?,
 
   creator_ip_address = ?
@@ -75,6 +85,8 @@ SET
     prompt_token.as_str(),
     args.prompt_type.to_str(),
     args.maybe_creator_user_token.map(|t| t.as_str()),
+    args.maybe_model_type.map(|m| m.to_str()),
+    args.maybe_generation_provider.map(|g| g.to_str()),
     args.maybe_positive_prompt,
     args.maybe_negative_prompt,
     maybe_other_args,
diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs
index 44d9d4ea9d..f8823ec276 100644
--- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs
+++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs
@@ -15,6 +15,8 @@ use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginM
 use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory;
 use enums::by_table::media_files::media_file_type::MediaFileType;
 use enums::by_table::prompts::prompt_type::PromptType;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
 use filesys::path_to_string::path_to_string;
 use mysql_queries::payloads::media_file_extra_info::inner_payloads::stable_diffusion_extra_info::StableDiffusionExtraInfo;
 use mysql_queries::payloads::media_file_extra_info::media_file_extra_info::MediaFileExtraInfo;
@@ -328,6 +330,8 @@ pub async fn process_job_inference(
     maybe_apriori_prompt_token: Some(&prompt_token),
     prompt_type: PromptType::StableDiffusion,
     maybe_creator_user_token: job.maybe_creator_user_token_typed.as_ref(),
+    maybe_model_type: None,
+    maybe_generation_provider: None,
     maybe_positive_prompt: Some(&positive_prompt),
     maybe_negative_prompt: maybe_negative_prompt.as_deref(),
     maybe_other_args: None, // TODO(bt,2024-02-22): Support other arguments
diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/validate_and_save_results.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/validate_and_save_results.rs
index e0a5152914..1ba6d2c4c6 100644
--- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/validate_and_save_results.rs
+++ b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/validate_and_save_results.rs
@@ -8,6 +8,8 @@ use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path:
 use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType;
 use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory;
 use enums::by_table::prompts::prompt_type::PromptType;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
 use filesys::file_size::file_size;
 use filesys::path_to_string::path_to_string;
 use hashing::sha256::sha256_hash_file::sha256_hash_file;
@@ -290,6 +292,8 @@ pub async fn validate_and_save_results(args: SaveResultsArgs<'_>) -> Result<Medi
       maybe_apriori_prompt_token: Some(&prompt_token),
       prompt_type: PromptType::ComfyUi,
       maybe_creator_user_token: args.job.maybe_creator_user_token_typed.as_ref(),
+      maybe_model_type: None,
+      maybe_generation_provider: None,
       maybe_positive_prompt: args.comfy_args.positive_prompt.as_deref(),
       maybe_negative_prompt: args.comfy_args.negative_prompt.as_deref(),
       maybe_other_args: maybe_other_args.as_ref(),
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gpt_image_1_edit_image_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gpt_image_1_edit_image_handler.rs
index 6766c3eeec..4c66f216f0 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gpt_image_1_edit_image_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/edit/gpt_image_1_edit_image_handler.rs
@@ -10,6 +10,8 @@ use actix_web::{web, HttpRequest};
 use artcraft_api_defs::generate::image::edit::gpt_image_1_edit_image::{GptImage1EditImageImageQuality, GptImage1EditImageImageSize, GptImage1EditImageNumImages, GptImage1EditImageRequest, GptImage1EditImageResponse};
 use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
 use enums::by_table::prompts::prompt_type::PromptType;
+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::webhook::image::enqueue_gpt_image_1_edit_image_webhook::enqueue_gpt_image_1_edit_image_webhook;
@@ -206,6 +208,8 @@ pub async fn gpt_image_1_edit_image_handler(
     maybe_creator_user_token: maybe_user_session
         .as_ref()
         .map(|s| &s.user_token),
+    maybe_model_type: Some(ModelType::GptImage1),
+    maybe_generation_provider: Some(GenerationProvider::Artcraft),
     maybe_positive_prompt: request.prompt.as_deref(),
     maybe_negative_prompt: None,
     maybe_other_args: None,
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/generate_veo_2_image_to_video_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/generate_veo_2_image_to_video_handler.rs
index 3cfcb87b19..a33e439718 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/generate_veo_2_image_to_video_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/generate_veo_2_image_to_video_handler.rs
@@ -15,6 +15,8 @@ use artcraft_api_defs::generate::video::generate_veo_2_image_to_video::GenerateV
 use artcraft_api_defs::generate::video::generate_veo_2_image_to_video::GenerateVeo2ImageToVideoResponse;
 use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
 use enums::by_table::prompts::prompt_type::PromptType;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
 use enums::common::visibility::Visibility;
 use fal_client::requests::webhook::video::enqueue_veo_2_image_to_video_webhook::enqueue_veo_2_image_to_video_webhook;
 use fal_client::requests::webhook::video::enqueue_veo_2_image_to_video_webhook::Veo2Args;
@@ -198,6 +200,8 @@ pub async fn generate_veo_2_image_to_video_handler(
     maybe_creator_user_token: maybe_user_session
         .as_ref()
         .map(|s| &s.user_token),
+    maybe_model_type: Some(ModelType::Veo2),
+    maybe_generation_provider: Some(GenerationProvider::Artcraft),
     maybe_positive_prompt: Some(prompt),
     maybe_negative_prompt: None,
     maybe_other_args: None,