Skip to main content

gproxy_protocol/protocol/openai/images/
requests.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize, de, de::DeserializeOwned};
4use serde_json::{Map, Value};
5
6use super::super::common::*;
7
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
9#[non_exhaustive]
10pub struct ImageGenerationRequest {
11    pub prompt: String,
12    #[serde(skip_serializing_if = "Option::is_none")]
13    pub background: Option<ImageBackground>,
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub model: Option<OpenAiModelId>,
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub moderation: Option<ImageModeration>,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub n: Option<u32>,
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub output_compression: Option<u32>,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub output_format: Option<ImageOutputFormat>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub partial_images: Option<u32>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub quality: Option<ImageQuality>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub response_format: Option<ImageResponseFormat>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub size: Option<ImageSize>,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub stream: Option<bool>,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub style: Option<ImageStyle>,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub user: Option<String>,
38    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
39    pub extra: Extra,
40}
41
42#[derive(Debug, Clone, PartialEq, Serialize, gproxy_protocol_macros::WireBuilder)]
43#[non_exhaustive]
44pub struct ImageEditRequest {
45    pub images: Vec<ImageReference>,
46    pub prompt: String,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub background: Option<ImageBackground>,
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub input_fidelity: Option<ImageInputFidelity>,
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub mask: Option<ImageReference>,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub model: Option<OpenAiModelId>,
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub moderation: Option<ImageModeration>,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub n: Option<u32>,
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub output_compression: Option<u32>,
61    #[serde(skip_serializing_if = "Option::is_none")]
62    pub output_format: Option<ImageOutputFormat>,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub partial_images: Option<u32>,
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub quality: Option<ImageEditQuality>,
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub size: Option<ImageEditSize>,
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub stream: Option<bool>,
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub user: Option<String>,
73    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
74    pub extra: Extra,
75}
76
77impl<'de> Deserialize<'de> for ImageEditRequest {
78    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
79    where
80        D: serde::Deserializer<'de>,
81    {
82        let value = Value::deserialize(deserializer)?;
83        let Value::Object(mut map) = value else {
84            return Err(de::Error::custom("image edit request must be an object"));
85        };
86
87        Ok(Self {
88            images: take_image_references(&mut map).map_err(de::Error::custom)?,
89            prompt: take_required(&mut map, "prompt").map_err(de::Error::custom)?,
90            background: take_optional(&mut map, "background").map_err(de::Error::custom)?,
91            input_fidelity: take_optional(&mut map, "input_fidelity").map_err(de::Error::custom)?,
92            mask: take_optional_image_reference(&mut map, "mask").map_err(de::Error::custom)?,
93            model: take_optional(&mut map, "model").map_err(de::Error::custom)?,
94            moderation: take_optional(&mut map, "moderation").map_err(de::Error::custom)?,
95            n: take_optional_u32(&mut map, "n").map_err(de::Error::custom)?,
96            output_compression: take_optional_u32(&mut map, "output_compression")
97                .map_err(de::Error::custom)?,
98            output_format: take_optional(&mut map, "output_format").map_err(de::Error::custom)?,
99            partial_images: take_optional_u32(&mut map, "partial_images")
100                .map_err(de::Error::custom)?,
101            quality: take_optional(&mut map, "quality").map_err(de::Error::custom)?,
102            size: take_optional(&mut map, "size").map_err(de::Error::custom)?,
103            stream: take_optional_bool(&mut map, "stream").map_err(de::Error::custom)?,
104            user: take_optional(&mut map, "user").map_err(de::Error::custom)?,
105            extra: map.into_iter().collect(),
106        })
107    }
108}
109
110#[derive(Debug, Clone, PartialEq, Serialize, gproxy_protocol_macros::WireBuilder)]
111#[non_exhaustive]
112pub struct ImageReference {
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub file_id: Option<String>,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub image_url: Option<String>,
117    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
118    pub extra: Extra,
119}
120
121impl<'de> Deserialize<'de> for ImageReference {
122    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
123    where
124        D: serde::Deserializer<'de>,
125    {
126        let value = Value::deserialize(deserializer)?;
127        if let Value::String(value) = value {
128            return string_image_reference(value).map_err(de::Error::custom);
129        }
130
131        #[derive(Deserialize)]
132        struct RawImageReference {
133            file_id: Option<String>,
134            image_url: Option<String>,
135            #[serde(default, flatten)]
136            extra: Extra,
137        }
138
139        let raw: RawImageReference = serde_json::from_value(value).map_err(de::Error::custom)?;
140        match (raw.file_id.is_some(), raw.image_url.is_some()) {
141            (true, false) | (false, true) => Ok(Self {
142                file_id: raw.file_id,
143                image_url: raw.image_url,
144                extra: raw.extra,
145            }),
146            (true, true) => Err(de::Error::custom(
147                "image reference must contain exactly one of file_id or image_url",
148            )),
149            (false, false) => Err(de::Error::custom(
150                "image reference must contain file_id or image_url",
151            )),
152        }
153    }
154}
155
156fn take_required<T: DeserializeOwned>(
157    map: &mut Map<String, Value>,
158    key: &str,
159) -> Result<T, String> {
160    let Some(value) = map.remove(key) else {
161        return Err(format!("missing required field `{key}`"));
162    };
163    serde_json::from_value(value).map_err(|e| format!("{key}: {e}"))
164}
165
166fn take_optional<T: DeserializeOwned>(
167    map: &mut Map<String, Value>,
168    key: &str,
169) -> Result<Option<T>, String> {
170    match map.remove(key) {
171        Some(Value::Null) | None => Ok(None),
172        Some(value) => serde_json::from_value(value)
173            .map(Some)
174            .map_err(|e| format!("{key}: {e}")),
175    }
176}
177
178fn take_optional_u32(map: &mut Map<String, Value>, key: &str) -> Result<Option<u32>, String> {
179    match map.remove(key) {
180        Some(Value::Null) | None => Ok(None),
181        Some(Value::String(value)) => value
182            .parse::<u32>()
183            .map(Some)
184            .map_err(|e| format!("{key}: {e}")),
185        Some(value) => serde_json::from_value(value)
186            .map(Some)
187            .map_err(|e| format!("{key}: {e}")),
188    }
189}
190
191fn take_optional_bool(map: &mut Map<String, Value>, key: &str) -> Result<Option<bool>, String> {
192    match map.remove(key) {
193        Some(Value::Null) | None => Ok(None),
194        Some(Value::String(value)) => value
195            .parse::<bool>()
196            .map(Some)
197            .map_err(|e| format!("{key}: {e}")),
198        Some(value) => serde_json::from_value(value)
199            .map(Some)
200            .map_err(|e| format!("{key}: {e}")),
201    }
202}
203
204fn take_image_references(map: &mut Map<String, Value>) -> Result<Vec<ImageReference>, String> {
205    let mut images = Vec::new();
206    if let Some(value) = map.remove("image") {
207        images.extend(image_references_from_value(value)?);
208    }
209    if let Some(value) = map.remove("images") {
210        images.extend(image_references_from_value(value)?);
211    }
212    if images.is_empty() {
213        return Err("missing required field `images`".to_owned());
214    }
215    Ok(images)
216}
217
218fn take_optional_image_reference(
219    map: &mut Map<String, Value>,
220    key: &str,
221) -> Result<Option<ImageReference>, String> {
222    match map.remove(key) {
223        Some(Value::Null) | None => Ok(None),
224        Some(value) => serde_json::from_value(value)
225            .map(Some)
226            .map_err(|e| format!("{key}: {e}")),
227    }
228}
229
230fn image_references_from_value(value: Value) -> Result<Vec<ImageReference>, String> {
231    match value {
232        Value::Array(values) => values
233            .into_iter()
234            .map(|value| serde_json::from_value(value).map_err(|e| e.to_string()))
235            .collect(),
236        value => serde_json::from_value(value)
237            .map(|reference| vec![reference])
238            .map_err(|e| e.to_string()),
239    }
240}
241
242fn string_image_reference(value: String) -> Result<ImageReference, String> {
243    if value.trim().is_empty() {
244        return Err("image reference string must not be empty".to_owned());
245    }
246    if value.starts_with("http://") || value.starts_with("https://") || value.starts_with("data:") {
247        Ok(ImageReference {
248            file_id: None,
249            image_url: Some(value),
250            extra: Default::default(),
251        })
252    } else {
253        Ok(ImageReference {
254            file_id: Some(value),
255            image_url: None,
256            extra: Default::default(),
257        })
258    }
259}