1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use http_req::{
    request::{Method, Request},
    uri::Uri,
};
use serde::Serialize;
use std::fmt;
use urlencoding::encode;

use crate::Retry;

#[derive(Debug)]
pub enum ImageSize {
    S256,
    S512,
    S1024,
}

impl fmt::Display for ImageSize {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ImageSize::S256 => write!(f, "256x256"),
            ImageSize::S512 => write!(f, "512x512"),
            ImageSize::S1024 => write!(f, "1024x1024"),
        }
    }
}

impl Serialize for ImageSize {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            ImageSize::S256 => serializer.serialize_str("256x256"),
            ImageSize::S512 => serializer.serialize_str("512x512"),
            ImageSize::S1024 => serializer.serialize_str("1024x1024"),
        }
    }
}

/// Request struct for the image creation.
///
/// For more detail about parameters, please refer to
/// [OpenAI docs](https://platform.openai.com/docs/api-reference/images/create)
///
#[derive(Debug, Serialize)]
pub struct ImageRequest {
    /// A text description of the desired image(s). The maximum length is 1000 characters.
    pub prompt: String,
    ///The number of images to generate. Must be between 1 and 10.
    pub n: u8,
    /// The size of the generated images. Must be one of 256x256, 512x512, or 1024x1024.
    pub size: ImageSize,
}

impl crate::OpenAIFlows {
    /// Create image for the provided prompt and parameters.
    ///
    /// `params` is a [ImageRequest] object.
    ///
    /// If you haven't connected your OpenAI account with [Flows.network platform](https://flows.network),
    /// you will receive an error in the flow's building log or running log.
    ///
    /// ```rust,no_run
    ///  // Create an instance of OpenAIFlows.
    ///  let openai_flows = OpenAIFlows::new();
    ///  // Create an `ImageRequest` instance.
    ///  let image_request = ImageRequest {
    ///      prompt: "An abstract representation of a summer sunrise".to_string(),
    ///      n: 2,
    ///      size: ImageSize::new(512, 512),
    ///  };
    ///
    ///  // Use `create_image` to generate images and handle the result.
    ///  match openai_flows.create_image(image_request).await {
    ///      Ok(images) => {
    ///          for (i, image) in images.iter().enumerate() {
    ///              println!("Image {}: {}", i + 1, image);
    ///          }
    ///          Ok(())
    ///      }
    ///      Err(e) => Err(format!("Failed to generate images: {}", e)),
    ///  }
    /// ```
    pub async fn create_image(&self, params: ImageRequest) -> Result<Vec<String>, String> {
        self.keep_trying(|account| create_image_inner(account, &params))
    }
}

fn create_image_inner(account: &str, params: &ImageRequest) -> Retry<Vec<String>> {
    let flows_user = unsafe { crate::_get_flows_user() };

    let mut writer = Vec::new();
    let uri = format!(
        "{}/{}/create_image?account={}",
        crate::OPENAI_API_PREFIX.as_str(),
        flows_user,
        encode(account),
    );
    let uri = Uri::try_from(uri.as_str()).unwrap();
    let body = serde_json::to_vec(params).unwrap_or_default();
    match Request::new(&uri)
        .method(Method::POST)
        .header("Content-Type", "application/json")
        .header("Content-Length", &body.len())
        .body(&body)
        .send(&mut writer)
    {
        Ok(res) => {
            match res.status_code().is_success() {
                true => Retry::No(
                    serde_json::from_slice::<Vec<String>>(&writer)
                        .or(Err(String::from("Unexpected error"))),
                ),
                false => {
                    match res.status_code().into() {
                        409 | 429 | 503 => {
                            // 409 TryAgain 429 RateLimitError
                            // 503 ServiceUnavailable
                            Retry::Yes(String::from_utf8_lossy(&writer).into_owned())
                        }
                        _ => Retry::No(Err(String::from_utf8_lossy(&writer).into_owned())),
                    }
                }
            }
        }
        Err(e) => Retry::No(Err(e.to_string())),
    }
}