zai_rs/model/gen_image/
data.rs

1use std::sync::Arc;
2
3use serde::Serialize;
4use validator::Validate;
5
6use super::{
7    super::traits::*,
8    image_request::{ImageGenBody, ImageQuality, ImageSize},
9};
10use crate::client::{
11    endpoints::{ApiBase, EndpointConfig, paths},
12    http::{HttpClient, HttpClientConfig, parse_typed_response},
13};
14
15/// Image generation request structure
16/// Provides a typed builder around the image generation API body
17pub struct ImageGenRequest<N>
18where
19    N: ModelName + ImageGen + Serialize,
20{
21    /// API key for authentication
22    pub key: String,
23    url: String,
24    endpoint_config: EndpointConfig,
25    api_base: ApiBase,
26    http_config: Arc<HttpClientConfig>,
27    /// Request body
28    body: ImageGenBody<N>,
29}
30
31impl<N> ImageGenRequest<N>
32where
33    N: ModelName + ImageGen + Serialize,
34{
35    /// Create a new image generation request for the given model and API key
36    pub fn new(model: N, key: String) -> Self {
37        let body = ImageGenBody {
38            model,
39            prompt: None,
40            quality: None,
41            size: None,
42            watermark_enabled: None,
43            user_id: None,
44        };
45        let endpoint_config = EndpointConfig::default();
46        let api_base = ApiBase::PaasV4;
47        let url = endpoint_config.url(&api_base, paths::IMAGES_GENERATIONS);
48        Self {
49            key,
50            url,
51            endpoint_config,
52            api_base,
53            http_config: Arc::new(HttpClientConfig::default()),
54            body,
55        }
56    }
57
58    fn rebuild_url(&mut self) {
59        self.url = self
60            .endpoint_config
61            .url(&self.api_base, paths::IMAGES_GENERATIONS);
62    }
63
64    pub fn with_base_url(mut self, base: impl Into<String>) -> Self {
65        self.api_base = ApiBase::Custom(base.into());
66        self.rebuild_url();
67        self
68    }
69
70    pub fn with_endpoint_config(mut self, endpoint_config: EndpointConfig) -> Self {
71        self.endpoint_config = endpoint_config;
72        self.rebuild_url();
73        self
74    }
75
76    pub fn with_http_config(mut self, config: HttpClientConfig) -> Self {
77        self.http_config = Arc::new(config);
78        self
79    }
80
81    /// Mutable access to inner body (for advanced customizations)
82    pub fn body_mut(&mut self) -> &mut ImageGenBody<N> {
83        &mut self.body
84    }
85
86    /// Set prompt text
87    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
88        self.body.prompt = Some(prompt.into());
89        self
90    }
91
92    /// Set image quality
93    pub fn with_quality(mut self, quality: ImageQuality) -> Self {
94        self.body.quality = Some(quality);
95        self
96    }
97
98    /// Set image size
99    pub fn with_size(mut self, size: ImageSize) -> Self {
100        self.body.size = Some(size);
101        self
102    }
103
104    /// Enable/disable watermark
105    pub fn with_watermark_enabled(mut self, watermark_enabled: bool) -> Self {
106        self.body.watermark_enabled = Some(watermark_enabled);
107        self
108    }
109
110    /// Set user id
111    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
112        self.body.user_id = Some(user_id.into());
113        self
114    }
115
116    pub fn validate(&self) -> crate::ZaiResult<()> {
117        // Body-level field validations
118        self.body
119            .validate()
120            .map_err(|e| crate::client::error::ZaiError::ApiError {
121                code: 1200,
122                message: format!("Validation error: {:?}", e),
123            })?;
124        // Require prompt
125        if self
126            .body
127            .prompt
128            .as_deref()
129            .map(|s| s.trim().is_empty())
130            .unwrap_or(true)
131        {
132            return Err(crate::client::error::ZaiError::ApiError {
133                code: 1200,
134                message: "prompt is required".to_string(),
135            });
136        }
137        // Validate custom size when present
138        if let Some(size) = &self.body.size
139            && let super::image_request::ImageSize::Custom { .. } = size
140            && !size.is_valid()
141        {
142            return Err(crate::client::error::ZaiError::ApiError {
143                        code: 1200,
144                        message: "invalid custom image size: must be 512..=2048, divisible by 16, and <= 2^21 pixels".to_string(),
145                    });
146        }
147        Ok(())
148    }
149
150    pub async fn send(&self) -> crate::ZaiResult<super::image_response::ImageResponse> {
151        self.validate()?;
152        let resp = self.post().await?;
153        let parsed = parse_typed_response::<super::image_response::ImageResponse>(resp).await?;
154        Ok(parsed)
155    }
156}
157
158impl<N> HttpClient for ImageGenRequest<N>
159where
160    N: ModelName + ImageGen + Serialize,
161{
162    type Body = ImageGenBody<N>;
163    type ApiUrl = String;
164    type ApiKey = String;
165
166    fn api_url(&self) -> &Self::ApiUrl {
167        &self.url
168    }
169
170    fn api_key(&self) -> &Self::ApiKey {
171        &self.key
172    }
173
174    fn body(&self) -> &Self::Body {
175        &self.body
176    }
177
178    fn http_config(&self) -> Arc<HttpClientConfig> {
179        Arc::clone(&self.http_config)
180    }
181}
zai_rs/model/gen_image/data.rs

zai_rs/model/gen_image/
data.rs
