zai_rs/services/images/
request.rs1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
46pub enum AsyncImageQuality {
47 #[serde(rename = "hd")]
49 Hd,
50}
51
52#[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 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 pub fn with_quality(mut self, quality: AsyncImageQuality) -> Self {
96 self.quality = Some(quality);
97 self
98 }
99
100 pub fn with_size(mut self, size: impl Into<String>) -> Self {
102 self.size = Some(size.into());
103 self
104 }
105
106 pub fn with_watermark(mut self, enabled: bool) -> Self {
108 self.watermark_enabled = Some(enabled);
109 self
110 }
111
112 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 pub const fn model(&self) -> &'static str {
120 "glm-image"
121 }
122
123 pub fn prompt(&self) -> &str {
125 &self.prompt
126 }
127
128 pub const fn quality(&self) -> Option<AsyncImageQuality> {
130 self.quality
131 }
132
133 pub fn size(&self) -> Option<&str> {
135 self.size.as_deref()
136 }
137
138 pub const fn watermark_enabled(&self) -> Option<bool> {
140 self.watermark_enabled
141 }
142
143 pub fn user_id(&self) -> Option<&str> {
145 self.user_id.as_deref()
146 }
147
148 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 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}