Skip to main content

zai_rs/model/gen_image/
data.rs

1use serde::Serialize;
2use validator::Validate;
3
4use super::{
5    super::traits::*,
6    image_request::{ImageGenBody, ImageQuality, ImageSize},
7};
8use crate::client::http::HttpClient;
9
10/// Image generation request structure
11/// Provides a typed builder around the image generation API body
12pub struct ImageGenRequest<N>
13where
14    N: ModelName + ImageGen + Serialize,
15{
16    /// API key for authentication
17    pub key: String,
18    /// Request body
19    body: ImageGenBody<N>,
20}
21
22impl<N> ImageGenRequest<N>
23where
24    N: ModelName + ImageGen + Serialize,
25{
26    /// Create a new image generation request for the given model and API key
27    pub fn new(model: N, key: String) -> Self {
28        let body = ImageGenBody {
29            model,
30            prompt: None,
31            quality: None,
32            size: None,
33            watermark_enabled: None,
34            user_id: None,
35        };
36        Self { key, body }
37    }
38
39    /// Mutable access to inner body (for advanced customizations)
40    pub fn body_mut(&mut self) -> &mut ImageGenBody<N> {
41        &mut self.body
42    }
43
44    /// Set prompt text
45    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
46        self.body.prompt = Some(prompt.into());
47        self
48    }
49
50    /// Set image quality
51    pub fn with_quality(mut self, quality: ImageQuality) -> Self {
52        self.body.quality = Some(quality);
53        self
54    }
55
56    /// Set image size
57    pub fn with_size(mut self, size: ImageSize) -> Self {
58        self.body.size = Some(size);
59        self
60    }
61
62    /// Enable/disable watermark
63    pub fn with_watermark_enabled(mut self, watermark_enabled: bool) -> Self {
64        self.body.watermark_enabled = Some(watermark_enabled);
65        self
66    }
67
68    /// Set user id
69    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
70        self.body.user_id = Some(user_id.into());
71        self
72    }
73
74    pub fn validate(&self) -> crate::ZaiResult<()> {
75        // Body-level field validations
76        self.body
77            .validate()
78            .map_err(|e| crate::client::error::ZaiError::ApiError {
79                code: 1200,
80                message: format!("Validation error: {:?}", e),
81            })?;
82        // Require prompt
83        if self
84            .body
85            .prompt
86            .as_deref()
87            .map(|s| s.trim().is_empty())
88            .unwrap_or(true)
89        {
90            return Err(crate::client::error::ZaiError::ApiError {
91                code: 1200,
92                message: "prompt is required".to_string(),
93            });
94        }
95        // Validate custom size when present
96        if let Some(size) = &self.body.size
97            && let super::image_request::ImageSize::Custom { .. } = size
98            && !size.is_valid()
99        {
100            return Err(crate::client::error::ZaiError::ApiError {
101                        code: 1200,
102                        message: "invalid custom image size: must be 512..=2048, divisible by 16, and <= 2^21 pixels".to_string(),
103                    });
104        }
105        Ok(())
106    }
107
108    pub async fn send(&self) -> crate::ZaiResult<super::image_response::ImageResponse> {
109        self.validate()?;
110        let resp = self.post().await?;
111        let parsed = resp.json::<super::image_response::ImageResponse>().await?;
112        Ok(parsed)
113    }
114}
115
116impl<N> HttpClient for ImageGenRequest<N>
117where
118    N: ModelName + ImageGen + Serialize,
119{
120    type Body = ImageGenBody<N>;
121    type ApiUrl = &'static str;
122    type ApiKey = String;
123
124    fn api_url(&self) -> &Self::ApiUrl {
125        &"https://open.bigmodel.cn/api/paas/v4/images/generations"
126    }
127
128    fn api_key(&self) -> &Self::ApiKey {
129        &self.key
130    }
131
132    fn body(&self) -> &Self::Body {
133        &self.body
134    }
135}