diff --git a/crates/lib/artcraft_router/src/client/mod.rs b/crates/lib/artcraft_router/src/client/mod.rs new file mode 100644 index 0000000000..4690f514ca --- /dev/null +++ b/crates/lib/artcraft_router/src/client/mod.rs @@ -0,0 +1,3 @@ +pub mod router_artcraft_client; +pub mod router_client; +pub mod router_client_builder; diff --git a/crates/lib/artcraft_router/src/client/router_artcraft_client.rs b/crates/lib/artcraft_router/src/client/router_artcraft_client.rs new file mode 100644 index 0000000000..09e04146e9 --- /dev/null +++ b/crates/lib/artcraft_router/src/client/router_artcraft_client.rs @@ -0,0 +1,7 @@ +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::utils::api_host::ApiHost; + +pub struct RouterArtcraftClient { + pub(crate) api_host: ApiHost, + pub(crate) credentials: StorytellerCredentialSet, +} diff --git a/crates/lib/artcraft_router/src/client/router_client.rs b/crates/lib/artcraft_router/src/client/router_client.rs new file mode 100644 index 0000000000..0a5f90c3e4 --- /dev/null +++ b/crates/lib/artcraft_router/src/client/router_client.rs @@ -0,0 +1,5 @@ +use crate::client::router_artcraft_client::RouterArtcraftClient; + +pub struct RouterClient { + pub(crate) artcraft_client: Option, +} diff --git a/crates/lib/artcraft_router/src/client/router_client_builder.rs b/crates/lib/artcraft_router/src/client/router_client_builder.rs new file mode 100644 index 0000000000..207c7f483d --- /dev/null +++ b/crates/lib/artcraft_router/src/client/router_client_builder.rs @@ -0,0 +1,25 @@ +use crate::client::router_artcraft_client::RouterArtcraftClient; +use crate::client::router_client::RouterClient; + +pub struct RouterClientBuilder { + artcraft_client: Option, +} + +impl RouterClientBuilder { + pub fn new() -> Self { + Self { + artcraft_client: None, + } + } + + pub fn set_artcraft_client(mut self, client: RouterArtcraftClient) -> Self { + self.artcraft_client = Some(client); + self + } + + pub fn build(self) -> RouterClient { + RouterClient { + artcraft_client: self.artcraft_client, + } + } +} diff --git a/crates/lib/artcraft_router/src/generate/generate_video/generate_video.rs b/crates/lib/artcraft_router/src/generate/generate_video/generate_video.rs index f3d1c0a4dd..ac2f4ce3f9 100644 --- a/crates/lib/artcraft_router/src/generate/generate_video/generate_video.rs +++ b/crates/lib/artcraft_router/src/generate/generate_video/generate_video.rs @@ -1,13 +1,13 @@ use crate::api::common_aspect_ratio::CommonAspectRatio; use crate::api::common_resolution::CommonVideoResolution; use crate::api::common_video_model::CommonVideoModel; +use crate::client::router_client::RouterClient; use crate::errors::artcraft_router_error::ArtcraftRouterError; use crate::generate::generate_video::providers::artcraft::generate_video_artcraft_seedance2p0::generate_video_artcraft_seedance2p0; -use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; -use storyteller_client::utils::api_host::ApiHost; use tokens::tokens::generic_inference_jobs::InferenceJobToken; -pub struct GenerateVideoArgs { +pub struct GenerateVideoArgs<'a> { + pub client: &'a RouterClient, pub model: CommonVideoModel, pub resolution: Option, pub aspect_ratio: Option, @@ -19,18 +19,12 @@ pub struct GenerateVideoResponse { pub all_inference_job_tokens: Vec, } -pub struct GenerateVideoContext<'a> { - pub api_host: &'a ApiHost, - pub storyteller_creds: Option<&'a StorytellerCredentialSet>, -} - pub async fn generate_video( - args: &GenerateVideoArgs, - context: &GenerateVideoContext<'_>, + args: &GenerateVideoArgs<'_>, ) -> Result { match args.model { CommonVideoModel::Seedance2p0 => { - generate_video_artcraft_seedance2p0(args, context.api_host, context.storyteller_creds).await + generate_video_artcraft_seedance2p0(args).await } _ => Err(ArtcraftRouterError::UnsupportedModel(format!("{:?}", args.model))), } diff --git a/crates/lib/artcraft_router/src/generate/generate_video/providers/artcraft/generate_video_artcraft_seedance2p0.rs b/crates/lib/artcraft_router/src/generate/generate_video/providers/artcraft/generate_video_artcraft_seedance2p0.rs index 0029fb296f..bd86248081 100644 --- a/crates/lib/artcraft_router/src/generate/generate_video/providers/artcraft/generate_video_artcraft_seedance2p0.rs +++ b/crates/lib/artcraft_router/src/generate/generate_video/providers/artcraft/generate_video_artcraft_seedance2p0.rs @@ -5,9 +5,7 @@ use crate::generate::generate_video::generate_video::{GenerateVideoArgs, Generat use artcraft_api_defs::generate::video::multi_function::seedance_2p0_multi_function_video_gen::{ Seedance2p0AspectRatio, Seedance2p0BatchCount, Seedance2p0MultiFunctionVideoGenRequest, }; -use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; use storyteller_client::endpoints::generate::video::multi_function::seedance_2p0_multi_function_video_gen::seedance_2p0_multi_function_video_gen; -use storyteller_client::utils::api_host::ApiHost; use uuid::Uuid; fn map_aspect_ratio(aspect_ratio: Option) -> Option { @@ -26,12 +24,14 @@ fn map_aspect_ratio(aspect_ratio: Option) -> Option, + args: &GenerateVideoArgs<'_>, ) -> Result { - let uuid_idempotency_token = Uuid::new_v4().to_string(); // TODO: Make this an optional top-level arg + let artcraft_client = args.client.artcraft_client.as_ref() + .ok_or_else(|| ArtcraftRouterError::InvalidInput( + "Artcraft client is not configured on the RouterClient".to_string() + ))?; + let uuid_idempotency_token = Uuid::new_v4().to_string(); // TODO: Make this an optional top-level arg let aspect_ratio = map_aspect_ratio(args.aspect_ratio); let request = Seedance2p0MultiFunctionVideoGenRequest { @@ -45,7 +45,11 @@ pub async fn generate_video_artcraft_seedance2p0( batch_count: Some(Seedance2p0BatchCount::One), }; - let response = seedance_2p0_multi_function_video_gen(api_host, creds, request) + let response = seedance_2p0_multi_function_video_gen( + &artcraft_client.api_host, + Some(&artcraft_client.credentials), + request, + ) .await .map_err(|err| ArtcraftRouterError::Provider(ProviderError::Storyteller(err)))?; diff --git a/crates/lib/artcraft_router/src/lib.rs b/crates/lib/artcraft_router/src/lib.rs index 8e15b974d3..94f7a8da83 100644 --- a/crates/lib/artcraft_router/src/lib.rs +++ b/crates/lib/artcraft_router/src/lib.rs @@ -1,3 +1,4 @@ pub mod api; +pub mod client; pub mod errors; pub mod generate;