Skip to main content

zai_rs/services/images/
request.rs

1//! Typed request for asynchronous GLM-Image generation.
2
3use serde::{Deserialize, Serialize};
4
5use crate::{
6    ZaiResult,
7    client::{
8        ZaiClient,
9        error::{ZaiError, codes},
10    },
11    model::AsyncResponse,
12};
13
14fn validation_error(message: impl Into<String>) -> ZaiError {
15    ZaiError::ApiError {
16        code: codes::SDK_VALIDATION,
17        message: message.into(),
18    }
19}
20
21fn valid_glm_image_size(value: &str) -> bool {
22    if matches!(value, "1728x960" | "960x1728") {
23        return true;
24    }
25    let Some((width, height)) = value.split_once('x') else {
26        return false;
27    };
28    let (Ok(width), Ok(height)) = (width.parse::<u32>(), height.parse::<u32>()) else {
29        return false;
30    };
31    (1024..=2048).contains(&width)
32        && (1024..=2048).contains(&height)
33        && width.is_multiple_of(32)
34        && height.is_multiple_of(32)
35        && u64::from(width) * u64::from(height) <= 1 << 22
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
39enum AsyncImageModel {
40    #[serde(rename = "glm-image")]
41    GlmImage,
42}
43
44/// Image quality accepted by the frozen async GLM-Image schema.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
46pub enum AsyncImageQuality {
47    /// High-definition generation.
48    #[serde(rename = "hd")]
49    Hd,
50}
51
52/// Request for `POST /async/images/generations`.
53#[derive(Clone, Serialize)]
54pub struct AsyncImageGenerationRequest {
55    model: AsyncImageModel,
56    prompt: String,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    quality: Option<AsyncImageQuality>,
59    #[serde(skip_serializing_if = "Option::is_none")]
60    size: Option<String>,
61    #[serde(skip_serializing_if = "Option::is_none")]
62    watermark_enabled: Option<bool>,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    user_id: Option<String>,
65}
66
67impl std::fmt::Debug for AsyncImageGenerationRequest {
68    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        formatter
70            .debug_struct("AsyncImageGenerationRequest")
71            .field("model", &self.model())
72            .field("prompt", &"[REDACTED]")
73            .field("quality", &self.quality)
74            .field("size", &self.size)
75            .field("watermark_enabled", &self.watermark_enabled)
76            .field("user_id", &self.user_id.as_ref().map(|_| "[REDACTED]"))
77            .finish()
78    }
79}
80
81impl AsyncImageGenerationRequest {
82    /// Create a GLM-Image request from the required prompt.
83    pub fn new(prompt: impl Into<String>) -> Self {
84        Self {
85            model: AsyncImageModel::GlmImage,
86            prompt: prompt.into(),
87            quality: None,
88            size: None,
89            watermark_enabled: None,
90            user_id: None,
91        }
92    }
93
94    /// Select the image quality.
95    pub fn with_quality(mut self, quality: AsyncImageQuality) -> Self {
96        self.quality = Some(quality);
97        self
98    }
99
100    /// Set a recommended or valid custom image size.
101    pub fn with_size(mut self, size: impl Into<String>) -> Self {
102        self.size = Some(size.into());
103        self
104    }
105
106    /// Configure generated-image watermarking.
107    pub fn with_watermark(mut self, enabled: bool) -> Self {
108        self.watermark_enabled = Some(enabled);
109        self
110    }
111
112    /// Set the end-user identifier used for abuse monitoring.
113    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
114        self.user_id = Some(user_id.into());
115        self
116    }
117
118    /// Return the fixed model identifier sent on the wire.
119    pub const fn model(&self) -> &'static str {
120        "glm-image"
121    }
122
123    /// Borrow the image prompt.
124    pub fn prompt(&self) -> &str {
125        &self.prompt
126    }
127
128    /// Return the configured quality.
129    pub const fn quality(&self) -> Option<AsyncImageQuality> {
130        self.quality
131    }
132
133    /// Borrow the configured image size.
134    pub fn size(&self) -> Option<&str> {
135        self.size.as_deref()
136    }
137
138    /// Return the configured watermark flag.
139    pub const fn watermark_enabled(&self) -> Option<bool> {
140        self.watermark_enabled
141    }
142
143    /// Borrow the configured end-user identifier.
144    pub fn user_id(&self) -> Option<&str> {
145        self.user_id.as_deref()
146    }
147
148    /// Validate documented prompt, size, and user-id constraints.
149    pub fn validate(&self) -> ZaiResult<()> {
150        if self.prompt.trim().is_empty() {
151            return Err(validation_error("async image prompt must not be blank"));
152        }
153        if self
154            .size
155            .as_deref()
156            .is_some_and(|size| !valid_glm_image_size(size))
157        {
158            return Err(validation_error("invalid async GLM-Image size"));
159        }
160        if self
161            .user_id
162            .as_ref()
163            .is_some_and(|value| !(6..=128).contains(&value.chars().count()))
164        {
165            return Err(validation_error(
166                "async image user_id must contain between 6 and 128 characters",
167            ));
168        }
169        Ok(())
170    }
171
172    /// Validate, send, and decode the async task response.
173    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<AsyncResponse> {
174        self.validate()?;
175        let route = crate::client::routes::IMAGES_GENERATE_ASYNC;
176        let url = client.endpoints().resolve_route(route, &[])?;
177        let response = client
178            .send_json::<_, AsyncResponse>(route.method(), url, self)
179            .await?;
180        response.validate()?;
181        Ok(response)
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn request_serializes_the_fixed_model_and_closed_fields() {
191        let request = AsyncImageGenerationRequest::new("a cat")
192            .with_quality(AsyncImageQuality::Hd)
193            .with_size("960x1728")
194            .with_watermark(true)
195            .with_user_id("user-1");
196        assert_eq!(
197            serde_json::to_value(request).unwrap(),
198            serde_json::json!({
199                "model": "glm-image",
200                "prompt": "a cat",
201                "quality": "hd",
202                "size": "960x1728",
203                "watermark_enabled": true,
204                "user_id": "user-1"
205            })
206        );
207        let debug = format!(
208            "{:?}",
209            AsyncImageGenerationRequest::new("private prompt").with_user_id("private-user")
210        );
211        assert!(!debug.contains("private prompt"));
212        assert!(!debug.contains("private-user"));
213    }
214
215    #[test]
216    fn validation_rejects_blank_prompts_and_invalid_documented_values() {
217        assert!(AsyncImageGenerationRequest::new(" ").validate().is_err());
218        assert!(
219            AsyncImageGenerationRequest::new("cat")
220                .with_size("1025x1024")
221                .validate()
222                .is_err()
223        );
224        assert!(
225            AsyncImageGenerationRequest::new("cat")
226                .with_user_id("short")
227                .validate()
228                .is_err()
229        );
230    }
231}