Skip to main content

origin_asset/asset/
mod.rs

1pub mod types;
2pub use types::*;
3
4use reqwest::Method;
5use serde::Deserialize;
6use serde_json::{json, Map, Value};
7
8use crate::error::Result;
9use crate::transport::HttpTransport;
10
11#[derive(Debug, Clone)]
12pub struct AssetClient {
13    pub(crate) transport: HttpTransport,
14}
15
16impl AssetClient {
17    pub fn new(transport: HttpTransport) -> Self {
18        Self { transport }
19    }
20
21    pub async fn generate(&self, request: &GenerateRequest) -> Result<GenerateResponse> {
22        let body = json!({
23            "asset_type": request.asset_type,
24            "prompt": request.prompt,
25            "model": request.model,
26            "input_file": request.input_file,
27            "provider": request.provider,
28            "size": request.size,
29            "transparent": request.transparent,
30            "reference_images": request.reference_images,
31            "edit_mode": request.edit_mode,
32            "session_id": request.session_id,
33            "params": request.params,
34        });
35
36        self.transport.post("/api/generate", &body).await
37    }
38
39    pub async fn generate_image(
40        &self,
41        prompt: impl Into<String>,
42        options: Option<ImageOptions>,
43    ) -> Result<GenerateResponse> {
44        let options = options.unwrap_or_default();
45        let request = GenerateRequest {
46            asset_type: AssetType::Image,
47            prompt: Some(prompt.into()),
48            model: options.model,
49            input_file: options.input,
50            provider: options.provider,
51            size: options.size,
52            transparent: options.transparent,
53            reference_images: options.reference_images,
54            edit_mode: options.edit_mode,
55            session_id: options.session_id,
56            params: options.params,
57        };
58
59        self.generate(&request).await
60    }
61
62    pub async fn generate_video(
63        &self,
64        prompt: impl Into<String>,
65        options: Option<VideoOptions>,
66    ) -> Result<GenerateResponse> {
67        let options = options.unwrap_or_default();
68        let request = GenerateRequest {
69            asset_type: AssetType::Video,
70            prompt: Some(prompt.into()),
71            model: options.model,
72            input_file: options.input,
73            provider: options.provider,
74            size: options.size,
75            transparent: None,
76            reference_images: Vec::new(),
77            edit_mode: None,
78            session_id: None,
79            params: options.params,
80        };
81
82        self.generate(&request).await
83    }
84
85    pub async fn generate_audio(
86        &self,
87        prompt: impl Into<String>,
88        options: Option<AudioOptions>,
89    ) -> Result<GenerateResponse> {
90        let options = options.unwrap_or_default();
91        let request = GenerateRequest {
92            asset_type: AssetType::Audio,
93            prompt: Some(prompt.into()),
94            model: options.model,
95            input_file: None,
96            provider: options.provider,
97            size: None,
98            transparent: None,
99            reference_images: Vec::new(),
100            edit_mode: None,
101            session_id: None,
102            params: merge_params(
103                options.params,
104                [
105                    option_entry("audio_type", options.audio_type),
106                    option_entry("duration_seconds", options.duration.map(Value::from)),
107                ],
108            ),
109        };
110
111        self.generate(&request).await
112    }
113
114    pub async fn generate_tts(
115        &self,
116        prompt: impl Into<String>,
117        options: Option<TtsOptions>,
118    ) -> Result<GenerateResponse> {
119        let options = options.unwrap_or_default();
120        let request = GenerateRequest {
121            asset_type: AssetType::Tts,
122            prompt: Some(prompt.into()),
123            model: options.model,
124            input_file: None,
125            provider: options.provider,
126            size: None,
127            transparent: None,
128            reference_images: Vec::new(),
129            edit_mode: None,
130            session_id: None,
131            params: merge_params(
132                options.params,
133                [
134                    option_entry("voice", options.voice),
135                    option_entry("voice_id", options.voice_id),
136                    option_entry("language_type", options.language),
137                    option_entry("instructions", options.instructions),
138                ],
139            ),
140        };
141
142        self.generate(&request).await
143    }
144
145    pub async fn generate_music(
146        &self,
147        prompt: impl Into<String>,
148        options: Option<MusicOptions>,
149    ) -> Result<GenerateResponse> {
150        let options = options.unwrap_or_default();
151        let request = GenerateRequest {
152            asset_type: AssetType::Music,
153            prompt: Some(prompt.into()),
154            model: options.model,
155            input_file: None,
156            provider: options.provider,
157            size: None,
158            transparent: None,
159            reference_images: Vec::new(),
160            edit_mode: None,
161            session_id: None,
162            params: merge_params(
163                options.params,
164                [
165                    option_entry("duration_seconds", options.duration.map(Value::from)),
166                    bool_entry("force_instrumental", options.force_instrumental),
167                    option_entry("output_format", options.output_format),
168                ],
169            ),
170        };
171
172        self.generate(&request).await
173    }
174
175    pub async fn generate_model3d(
176        &self,
177        prompt: impl Into<String>,
178        options: Option<Model3dOptions>,
179    ) -> Result<GenerateResponse> {
180        let options = options.unwrap_or_default();
181        let request = GenerateRequest {
182            asset_type: AssetType::Model3d,
183            prompt: Some(prompt.into()),
184            model: options.model,
185            input_file: options.input,
186            provider: options.provider,
187            size: None,
188            transparent: None,
189            reference_images: Vec::new(),
190            edit_mode: None,
191            session_id: None,
192            params: merge_params(
193                options.params,
194                [
195                    option_entry("model_version", options.model_version),
196                    option_entry("face_limit", options.face_limit.map(Value::from)),
197                    bool_entry("pbr", options.pbr),
198                    option_entry("texture_quality", options.texture_quality),
199                    bool_entry("auto_size", options.auto_size),
200                    option_entry("negative_prompt", options.negative_prompt),
201                    option_entry(
202                        "multiview",
203                        (!options.multiview.is_empty()).then(|| {
204                            Value::Array(options.multiview.into_iter().map(Value::from).collect())
205                        }),
206                    ),
207                    option_entry("style", options.style),
208                ],
209            ),
210        };
211
212        self.generate(&request).await
213    }
214
215    pub async fn generate_sprite(
216        &self,
217        prompt: impl Into<String>,
218        options: Option<SpriteOptions>,
219    ) -> Result<GenerateResponse> {
220        let options = options.unwrap_or_default();
221        let request = GenerateRequest {
222            asset_type: AssetType::Sprite,
223            prompt: Some(prompt.into()),
224            model: options.model,
225            input_file: options.input,
226            provider: options.provider,
227            size: None,
228            transparent: None,
229            reference_images: Vec::new(),
230            edit_mode: None,
231            session_id: None,
232            params: merge_params(
233                options.params,
234                [
235                    option_entry(
236                        "animation_type",
237                        Some(Value::String(
238                            options.animation_type.unwrap_or_else(|| "walk".to_string()),
239                        )),
240                    ),
241                    option_entry(
242                        "direction",
243                        Some(Value::String(
244                            options.direction.unwrap_or_else(|| "right".to_string()),
245                        )),
246                    ),
247                    option_entry("duration", Some(Value::from(options.duration.unwrap_or(2)))),
248                    option_entry(
249                        "output_format",
250                        Some(Value::String(
251                            options
252                                .output_format
253                                .unwrap_or_else(|| "spritesheet".to_string()),
254                        )),
255                    ),
256                    option_entry("fps", options.fps.map(Value::from)),
257                    option_entry("style", options.style),
258                ],
259            ),
260        };
261
262        self.generate(&request).await
263    }
264
265    pub async fn process(&self, request: &ProcessRequest) -> Result<ProcessResponse> {
266        let body = json!({
267            "input": request.input,
268            "inputs": request.inputs,
269            "operations": request.operations,
270        });
271
272        self.transport.post("/api/process", &body).await
273    }
274
275    pub async fn jobs(&self, status: Option<&str>, limit: Option<u32>) -> Result<Vec<JobSummary>> {
276        let mut path = String::from("/api/jobs");
277        let mut query = Vec::new();
278
279        if let Some(status) = status {
280            query.push(format!("status={status}"));
281        }
282        if let Some(limit) = limit {
283            query.push(format!("limit={limit}"));
284        }
285        if !query.is_empty() {
286            path.push('?');
287            path.push_str(&query.join("&"));
288        }
289
290        let response: JobListResponse = self.transport.get(&path).await?;
291        Ok(response.jobs)
292    }
293
294    pub async fn job_status(&self, job_id: &str) -> Result<JobSummary> {
295        self.transport.get(&format!("/api/jobs/{job_id}")).await
296    }
297
298    pub async fn providers(&self) -> Result<Vec<ProviderInfo>> {
299        let response: ProviderListResponse = self.transport.get("/api/providers").await?;
300        Ok(response.providers)
301    }
302
303    pub async fn health(&self) -> Result<bool> {
304        let response: HealthResponse = self.transport.request(Method::GET, "/health", None).await?;
305        Ok(response.status == "healthy")
306    }
307}
308
309#[derive(Debug, Deserialize)]
310struct JobListResponse {
311    jobs: Vec<JobSummary>,
312}
313
314#[derive(Debug, Deserialize)]
315struct ProviderListResponse {
316    providers: Vec<ProviderInfo>,
317}
318
319#[derive(Debug, Deserialize)]
320struct HealthResponse {
321    status: String,
322}
323
324fn merge_params<const N: usize>(base: Value, entries: [(String, Option<Value>); N]) -> Value {
325    let mut params = match base {
326        Value::Object(map) => map,
327        _ => Map::new(),
328    };
329
330    for (key, value) in entries {
331        if let Some(value) = value {
332            params.insert(key, value);
333        }
334    }
335
336    if params.is_empty() {
337        Value::Null
338    } else {
339        Value::Object(params)
340    }
341}
342
343fn option_entry<T>(key: &str, value: Option<T>) -> (String, Option<Value>)
344where
345    T: Into<Value>,
346{
347    (key.to_string(), value.map(Into::into))
348}
349
350fn bool_entry(key: &str, value: Option<bool>) -> (String, Option<Value>) {
351    (
352        key.to_string(),
353        value.and_then(|enabled| enabled.then_some(Value::Bool(true))),
354    )
355}