zai_rs/model/gen_image/
data.rs1use serde::Serialize;
2use validator::Validate;
3
4use super::{
5 super::traits::*,
6 image_request::{ImageGenBody, ImageQuality, ImageSize},
7};
8use crate::client::http::HttpClient;
9
10pub struct ImageGenRequest<N>
13where
14 N: ModelName + ImageGen + Serialize,
15{
16 pub key: String,
18 body: ImageGenBody<N>,
20}
21
22impl<N> ImageGenRequest<N>
23where
24 N: ModelName + ImageGen + Serialize,
25{
26 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 pub fn body_mut(&mut self) -> &mut ImageGenBody<N> {
41 &mut self.body
42 }
43
44 pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
46 self.body.prompt = Some(prompt.into());
47 self
48 }
49
50 pub fn with_quality(mut self, quality: ImageQuality) -> Self {
52 self.body.quality = Some(quality);
53 self
54 }
55
56 pub fn with_size(mut self, size: ImageSize) -> Self {
58 self.body.size = Some(size);
59 self
60 }
61
62 pub fn with_watermark_enabled(mut self, watermark_enabled: bool) -> Self {
64 self.body.watermark_enabled = Some(watermark_enabled);
65 self
66 }
67
68 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 self.body
77 .validate()
78 .map_err(|e| crate::client::error::ZaiError::ApiError {
79 code: 1200,
80 message: format!("Validation error: {:?}", e),
81 })?;
82 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 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}