portkey_sdk/service/images.rs
1//! Images API service.
2//!
3//! This module provides methods for generating images using DALL-E models.
4
5use std::future::Future;
6
7use reqwest::multipart::{Form, Part};
8
9use crate::client::PortkeyClient;
10use crate::error::Result;
11use crate::model::{
12 CreateImageEditRequest, CreateImageRequest, CreateImageVariationRequest, ImageSize,
13 ImagesResponse,
14};
15
16/// Trait for Images API operations.
17pub trait ImagesService {
18 /// Generates images based on a text prompt.
19 ///
20 /// # Arguments
21 ///
22 /// * `request` - The image generation request containing the prompt and parameters
23 ///
24 /// # Returns
25 ///
26 /// Returns an `ImagesResponse` containing the generated images.
27 ///
28 /// # Example
29 ///
30 /// ```no_run
31 /// use portkey_sdk::{PortkeyConfig, PortkeyClient};
32 /// use portkey_sdk::service::ImagesService;
33 /// use portkey_sdk::builder::AuthMethod;
34 /// use portkey_sdk::model::{CreateImageRequest, ImageSize};
35 ///
36 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
37 /// let config = PortkeyConfig::builder()
38 /// .with_api_key("your-api-key")
39 /// .with_auth_method(AuthMethod::virtual_key("your-virtual-key"))
40 /// .build()?;
41 /// let client = PortkeyClient::new(config)?;
42 ///
43 /// let request = CreateImageRequest {
44 /// prompt: "A cute baby sea otter".to_string(),
45 /// model: Some("dall-e-3".to_string()),
46 /// n: Some(1),
47 /// size: Some(ImageSize::Size1024x1024),
48 /// response_format: None,
49 /// quality: None,
50 /// style: None,
51 /// user: None,
52 /// };
53 ///
54 /// let response = client.generate_image(request).await?;
55 /// println!("Generated {} images", response.data.len());
56 /// # Ok(())
57 /// # }
58 /// ```
59 fn generate_image(
60 &self,
61 request: CreateImageRequest,
62 ) -> impl Future<Output = Result<ImagesResponse>>;
63
64 /// Edits an image based on a prompt.
65 ///
66 /// # Arguments
67 ///
68 /// * `image_data` - The image file data as bytes (PNG, <4MB, square)
69 /// * `image_name` - The name of the image file
70 /// * `mask_data` - Optional mask image data (transparent areas indicate where to edit)
71 /// * `mask_name` - The name of the mask file (required if mask_data is provided)
72 /// * `request` - The image edit request containing the prompt and parameters
73 ///
74 /// # Returns
75 ///
76 /// Returns an `ImagesResponse` containing the edited images.
77 ///
78 /// # Example
79 ///
80 /// ```no_run
81 /// use portkey_sdk::{PortkeyConfig, PortkeyClient};
82 /// use portkey_sdk::service::ImagesService;
83 /// use portkey_sdk::builder::AuthMethod;
84 /// use portkey_sdk::model::{CreateImageEditRequest, ImageSize};
85 /// use std::fs;
86 ///
87 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
88 /// let config = PortkeyConfig::builder()
89 /// .with_api_key("your-api-key")
90 /// .with_auth_method(AuthMethod::virtual_key("your-virtual-key"))
91 /// .build()?;
92 /// let client = PortkeyClient::new(config)?;
93 ///
94 /// let image_data = fs::read("otter.png")?;
95 /// let mask_data = fs::read("mask.png")?;
96 ///
97 /// let request = CreateImageEditRequest {
98 /// prompt: "A cute baby sea otter wearing a beret".to_string(),
99 /// model: Some("dall-e-2".to_string()),
100 /// n: Some(1),
101 /// size: Some(ImageSize::Size1024x1024),
102 /// response_format: None,
103 /// user: None,
104 /// };
105 ///
106 /// let response = client.edit_image(
107 /// image_data,
108 /// "otter.png",
109 /// Some(mask_data),
110 /// Some("mask.png"),
111 /// request
112 /// ).await?;
113 /// # Ok(())
114 /// # }
115 /// ```
116 fn edit_image(
117 &self,
118 image_data: Vec<u8>,
119 image_name: &str,
120 mask_data: Option<Vec<u8>>,
121 mask_name: Option<&str>,
122 request: CreateImageEditRequest,
123 ) -> impl Future<Output = Result<ImagesResponse>>;
124
125 /// Creates a variation of an image.
126 ///
127 /// # Arguments
128 ///
129 /// * `image_data` - The image file data as bytes (PNG, <4MB, square)
130 /// * `image_name` - The name of the image file
131 /// * `request` - The image variation request containing parameters
132 ///
133 /// # Returns
134 ///
135 /// Returns an `ImagesResponse` containing the image variations.
136 ///
137 /// # Example
138 ///
139 /// ```no_run
140 /// use portkey_sdk::{PortkeyConfig, PortkeyClient};
141 /// use portkey_sdk::service::ImagesService;
142 /// use portkey_sdk::builder::AuthMethod;
143 /// use portkey_sdk::model::{CreateImageVariationRequest, ImageSize};
144 /// use std::fs;
145 ///
146 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
147 /// let config = PortkeyConfig::builder()
148 /// .with_api_key("your-api-key")
149 /// .with_auth_method(AuthMethod::virtual_key("your-virtual-key"))
150 /// .build()?;
151 /// let client = PortkeyClient::new(config)?;
152 ///
153 /// let image_data = fs::read("otter.png")?;
154 ///
155 /// let request = CreateImageVariationRequest {
156 /// model: Some("dall-e-2".to_string()),
157 /// n: Some(2),
158 /// size: Some(ImageSize::Size1024x1024),
159 /// ..Default::default()
160 /// };
161 ///
162 /// let response = client.create_image_variation(
163 /// image_data,
164 /// "otter.png",
165 /// request
166 /// ).await?;
167 /// # Ok(())
168 /// # }
169 /// ```
170 fn create_image_variation(
171 &self,
172 image_data: Vec<u8>,
173 image_name: &str,
174 request: CreateImageVariationRequest,
175 ) -> impl Future<Output = Result<ImagesResponse>>;
176}
177
178impl ImagesService for PortkeyClient {
179 async fn generate_image(&self, request: CreateImageRequest) -> Result<ImagesResponse> {
180 let response = self
181 .send_json(reqwest::Method::POST, "/images/generations", &request)
182 .await?;
183 let response = response.error_for_status()?;
184 let images_response: ImagesResponse = response.json().await?;
185 Ok(images_response)
186 }
187
188 async fn edit_image(
189 &self,
190 image_data: Vec<u8>,
191 image_name: &str,
192 mask_data: Option<Vec<u8>>,
193 mask_name: Option<&str>,
194 request: CreateImageEditRequest,
195 ) -> Result<ImagesResponse> {
196 // Build multipart form
197 let image_part = Part::bytes(image_data).file_name(image_name.to_string());
198
199 let mut form = Form::new()
200 .part("image", image_part)
201 .text("prompt", request.prompt.clone());
202
203 // Add optional mask
204 if let (Some(mask_bytes), Some(mask_filename)) = (mask_data, mask_name) {
205 let mask_part = Part::bytes(mask_bytes).file_name(mask_filename.to_string());
206 form = form.part("mask", mask_part);
207 }
208
209 // Add optional parameters
210 if let Some(model) = request.model {
211 form = form.text("model", model);
212 }
213
214 if let Some(n) = request.n {
215 form = form.text("n", n.to_string());
216 }
217
218 if let Some(size) = request.size {
219 let size_str = match size {
220 ImageSize::Size256x256 => "256x256",
221 ImageSize::Size512x512 => "512x512",
222 ImageSize::Size1024x1024 => "1024x1024",
223 ImageSize::Size1792x1024 => "1792x1024",
224 ImageSize::Size1024x1792 => "1024x1792",
225 };
226 form = form.text("size", size_str);
227 }
228
229 if let Some(response_format) = request.response_format {
230 let format_str = match response_format {
231 crate::model::ImageResponseFormat::Url => "url",
232 crate::model::ImageResponseFormat::B64Json => "b64_json",
233 };
234 form = form.text("response_format", format_str);
235 }
236
237 if let Some(user) = request.user {
238 form = form.text("user", user);
239 }
240
241 let response = self
242 .send_multipart(reqwest::Method::POST, "/images/edits", form)
243 .await?;
244
245 let response = response.error_for_status()?;
246 let images_response: ImagesResponse = response.json().await?;
247 Ok(images_response)
248 }
249
250 async fn create_image_variation(
251 &self,
252 image_data: Vec<u8>,
253 image_name: &str,
254 request: CreateImageVariationRequest,
255 ) -> Result<ImagesResponse> {
256 // Build multipart form
257 let image_part = Part::bytes(image_data).file_name(image_name.to_string());
258
259 let mut form = Form::new().part("image", image_part);
260
261 // Add optional parameters
262 if let Some(model) = request.model {
263 form = form.text("model", model);
264 }
265
266 if let Some(n) = request.n {
267 form = form.text("n", n.to_string());
268 }
269
270 if let Some(size) = request.size {
271 let size_str = match size {
272 ImageSize::Size256x256 => "256x256",
273 ImageSize::Size512x512 => "512x512",
274 ImageSize::Size1024x1024 => "1024x1024",
275 ImageSize::Size1792x1024 => "1792x1024",
276 ImageSize::Size1024x1792 => "1024x1792",
277 };
278 form = form.text("size", size_str);
279 }
280
281 if let Some(response_format) = request.response_format {
282 let format_str = match response_format {
283 crate::model::ImageResponseFormat::Url => "url",
284 crate::model::ImageResponseFormat::B64Json => "b64_json",
285 };
286 form = form.text("response_format", format_str);
287 }
288
289 if let Some(user) = request.user {
290 form = form.text("user", user);
291 }
292
293 let response = self
294 .send_multipart(reqwest::Method::POST, "/images/variations", form)
295 .await?;
296
297 let response = response.error_for_status()?;
298 let images_response: ImagesResponse = response.json().await?;
299 Ok(images_response)
300 }
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306 use crate::model::{ImageQuality, ImageSize, ImageStyle};
307
308 #[test]
309 fn test_create_image_request() {
310 let request = CreateImageRequest {
311 prompt: "A cute baby sea otter".to_string(),
312 model: Some("dall-e-3".to_string()),
313 n: Some(1),
314 quality: Some(ImageQuality::Hd),
315 size: Some(ImageSize::Size1024x1024),
316 style: Some(ImageStyle::Vivid),
317 response_format: None,
318 user: None,
319 };
320
321 assert_eq!(request.prompt, "A cute baby sea otter");
322 assert_eq!(request.model, Some("dall-e-3".to_string()));
323 assert_eq!(request.n, Some(1));
324 }
325}