Skip to main content

zai_rs/model/gen_image/
data.rs

1use validator::Validate;
2
3use super::{
4    super::traits::*,
5    image_request::{ImageGenBody, ImageQuality, ImageSize},
6};
7use crate::client::ZaiClient;
8
9/// Typed builder for an image-generation request sent through a [`ZaiClient`].
10pub struct ImageGenRequest<N>
11where
12    N: ImageGen,
13{
14    /// Request body
15    body: ImageGenBody<N>,
16}
17
18impl<N> ImageGenRequest<N>
19where
20    N: ImageGen,
21{
22    /// Create a new image generation request for the given model.
23    pub fn new(model: N) -> Self {
24        Self {
25            body: ImageGenBody::new(model),
26        }
27    }
28
29    /// Mutable access to inner body (for advanced customizations)
30    pub fn body_mut(&mut self) -> &mut ImageGenBody<N> {
31        &mut self.body
32    }
33
34    /// Set prompt text
35    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
36        self.body.prompt = Some(prompt.into());
37        self
38    }
39
40    /// Set image quality
41    pub fn with_quality(mut self, quality: ImageQuality) -> Self {
42        self.body.quality = Some(quality);
43        self
44    }
45
46    /// Set image size
47    pub fn with_size(mut self, size: ImageSize) -> Self {
48        self.body.size = Some(size);
49        self
50    }
51
52    /// Enable/disable watermark
53    pub fn with_watermark_enabled(mut self, watermark_enabled: bool) -> Self {
54        self.body.watermark_enabled = Some(watermark_enabled);
55        self
56    }
57
58    /// Set user id
59    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
60        self.body.user_id = Some(user_id.into());
61        self
62    }
63
64    /// Validate body constraints: required prompt and (when set) a valid
65    /// custom image size.
66    pub fn validate(&self) -> crate::ZaiResult<()> {
67        self.body
68            .validate()
69            .map_err(crate::client::error::ZaiError::from)?;
70        Ok(())
71    }
72
73    /// Submit the request via a [`ZaiClient`] and parse the typed
74    /// image-generation response.
75    pub async fn send_via(
76        &self,
77        client: &ZaiClient,
78    ) -> crate::ZaiResult<super::image_response::ImageResponse> {
79        self.validate()?;
80        let route = crate::client::routes::IMAGES_GENERATE;
81        let url = client.endpoints().resolve_route(route, &[])?;
82        client
83            .send_json::<_, super::image_response::ImageResponse>(route.method(), url, &self.body)
84            .await
85    }
86}