gproxy_protocol/protocol/openai/
images.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize, de, de::DeserializeOwned};
4use serde_json::{Map, Value};
5
6use super::common::*;
7
8pub type ImageGenerationWireModel = OpenAiWireModel<ImageGenerationRequest, ImagesResponse>;
9pub type ImageGenerationStreamWireModel =
10    OpenAiWireModel<ImageGenerationRequest, ImageGenerationStreamEvent>;
11pub type ImageEditWireModel = OpenAiWireModel<ImageEditRequest, ImagesResponse>;
12pub type ImageEditStreamWireModel = OpenAiWireModel<ImageEditRequest, ImageEditStreamEvent>;
13
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct ImageGenerationRequest {
16    pub prompt: String,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub background: Option<ImageBackground>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub model: Option<OpenAiModelId>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub moderation: Option<ImageModeration>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub n: Option<u32>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub output_compression: Option<u32>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub output_format: Option<ImageOutputFormat>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub partial_images: Option<u32>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub quality: Option<ImageQuality>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub response_format: Option<ImageResponseFormat>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub size: Option<ImageSize>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub stream: Option<bool>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub style: Option<ImageStyle>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub user: Option<String>,
43    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
44    pub extra: Extra,
45}
46
47#[derive(Debug, Clone, PartialEq, Serialize)]
48pub struct ImageEditRequest {
49    pub images: Vec<ImageReference>,
50    pub prompt: String,
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub background: Option<ImageBackground>,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub input_fidelity: Option<ImageInputFidelity>,
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub mask: Option<ImageReference>,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub model: Option<OpenAiModelId>,
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub moderation: Option<ImageModeration>,
61    #[serde(skip_serializing_if = "Option::is_none")]
62    pub n: Option<u32>,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub output_compression: Option<u32>,
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub output_format: Option<ImageOutputFormat>,
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub partial_images: Option<u32>,
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub quality: Option<ImageEditQuality>,
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub size: Option<ImageEditSize>,
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub stream: Option<bool>,
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub user: Option<String>,
77    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
78    pub extra: Extra,
79}
80
81impl<'de> Deserialize<'de> for ImageEditRequest {
82    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
83    where
84        D: serde::Deserializer<'de>,
85    {
86        let value = Value::deserialize(deserializer)?;
87        let Value::Object(mut map) = value else {
88            return Err(de::Error::custom("image edit request must be an object"));
89        };
90
91        Ok(Self {
92            images: take_image_references(&mut map).map_err(de::Error::custom)?,
93            prompt: take_required(&mut map, "prompt").map_err(de::Error::custom)?,
94            background: take_optional(&mut map, "background").map_err(de::Error::custom)?,
95            input_fidelity: take_optional(&mut map, "input_fidelity").map_err(de::Error::custom)?,
96            mask: take_optional_image_reference(&mut map, "mask").map_err(de::Error::custom)?,
97            model: take_optional(&mut map, "model").map_err(de::Error::custom)?,
98            moderation: take_optional(&mut map, "moderation").map_err(de::Error::custom)?,
99            n: take_optional_u32(&mut map, "n").map_err(de::Error::custom)?,
100            output_compression: take_optional_u32(&mut map, "output_compression")
101                .map_err(de::Error::custom)?,
102            output_format: take_optional(&mut map, "output_format").map_err(de::Error::custom)?,
103            partial_images: take_optional_u32(&mut map, "partial_images")
104                .map_err(de::Error::custom)?,
105            quality: take_optional(&mut map, "quality").map_err(de::Error::custom)?,
106            size: take_optional(&mut map, "size").map_err(de::Error::custom)?,
107            stream: take_optional_bool(&mut map, "stream").map_err(de::Error::custom)?,
108            user: take_optional(&mut map, "user").map_err(de::Error::custom)?,
109            extra: map.into_iter().collect(),
110        })
111    }
112}
113
114#[derive(Debug, Clone, PartialEq, Serialize)]
115pub struct ImageReference {
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub file_id: Option<String>,
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub image_url: Option<String>,
120    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
121    pub extra: Extra,
122}
123
124impl<'de> Deserialize<'de> for ImageReference {
125    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
126    where
127        D: serde::Deserializer<'de>,
128    {
129        let value = Value::deserialize(deserializer)?;
130        if let Value::String(value) = value {
131            return string_image_reference(value).map_err(de::Error::custom);
132        }
133
134        #[derive(Deserialize)]
135        struct RawImageReference {
136            file_id: Option<String>,
137            image_url: Option<String>,
138            #[serde(default, flatten)]
139            extra: Extra,
140        }
141
142        let raw: RawImageReference = serde_json::from_value(value).map_err(de::Error::custom)?;
143        match (raw.file_id.is_some(), raw.image_url.is_some()) {
144            (true, false) | (false, true) => Ok(Self {
145                file_id: raw.file_id,
146                image_url: raw.image_url,
147                extra: raw.extra,
148            }),
149            (true, true) => Err(de::Error::custom(
150                "image reference must contain exactly one of file_id or image_url",
151            )),
152            (false, false) => Err(de::Error::custom(
153                "image reference must contain file_id or image_url",
154            )),
155        }
156    }
157}
158
159fn take_required<T: DeserializeOwned>(
160    map: &mut Map<String, Value>,
161    key: &str,
162) -> Result<T, String> {
163    let Some(value) = map.remove(key) else {
164        return Err(format!("missing required field `{key}`"));
165    };
166    serde_json::from_value(value).map_err(|e| format!("{key}: {e}"))
167}
168
169fn take_optional<T: DeserializeOwned>(
170    map: &mut Map<String, Value>,
171    key: &str,
172) -> Result<Option<T>, String> {
173    match map.remove(key) {
174        Some(Value::Null) | None => Ok(None),
175        Some(value) => serde_json::from_value(value)
176            .map(Some)
177            .map_err(|e| format!("{key}: {e}")),
178    }
179}
180
181fn take_optional_u32(map: &mut Map<String, Value>, key: &str) -> Result<Option<u32>, String> {
182    match map.remove(key) {
183        Some(Value::Null) | None => Ok(None),
184        Some(Value::String(value)) => value
185            .parse::<u32>()
186            .map(Some)
187            .map_err(|e| format!("{key}: {e}")),
188        Some(value) => serde_json::from_value(value)
189            .map(Some)
190            .map_err(|e| format!("{key}: {e}")),
191    }
192}
193
194fn take_optional_bool(map: &mut Map<String, Value>, key: &str) -> Result<Option<bool>, String> {
195    match map.remove(key) {
196        Some(Value::Null) | None => Ok(None),
197        Some(Value::String(value)) => value
198            .parse::<bool>()
199            .map(Some)
200            .map_err(|e| format!("{key}: {e}")),
201        Some(value) => serde_json::from_value(value)
202            .map(Some)
203            .map_err(|e| format!("{key}: {e}")),
204    }
205}
206
207fn take_image_references(map: &mut Map<String, Value>) -> Result<Vec<ImageReference>, String> {
208    let mut images = Vec::new();
209    if let Some(value) = map.remove("image") {
210        images.extend(image_references_from_value(value)?);
211    }
212    if let Some(value) = map.remove("images") {
213        images.extend(image_references_from_value(value)?);
214    }
215    if images.is_empty() {
216        return Err("missing required field `images`".to_owned());
217    }
218    Ok(images)
219}
220
221fn take_optional_image_reference(
222    map: &mut Map<String, Value>,
223    key: &str,
224) -> Result<Option<ImageReference>, String> {
225    match map.remove(key) {
226        Some(Value::Null) | None => Ok(None),
227        Some(value) => serde_json::from_value(value)
228            .map(Some)
229            .map_err(|e| format!("{key}: {e}")),
230    }
231}
232
233fn image_references_from_value(value: Value) -> Result<Vec<ImageReference>, String> {
234    match value {
235        Value::Array(values) => values
236            .into_iter()
237            .map(|value| serde_json::from_value(value).map_err(|e| e.to_string()))
238            .collect(),
239        value => serde_json::from_value(value)
240            .map(|reference| vec![reference])
241            .map_err(|e| e.to_string()),
242    }
243}
244
245fn string_image_reference(value: String) -> Result<ImageReference, String> {
246    if value.trim().is_empty() {
247        return Err("image reference string must not be empty".to_owned());
248    }
249    if value.starts_with("http://") || value.starts_with("https://") || value.starts_with("data:") {
250        Ok(ImageReference {
251            file_id: None,
252            image_url: Some(value),
253            extra: Default::default(),
254        })
255    } else {
256        Ok(ImageReference {
257            file_id: Some(value),
258            image_url: None,
259            extra: Default::default(),
260        })
261    }
262}
263
264#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
265pub struct ImagesResponse {
266    pub created: u64,
267    #[serde(skip_serializing_if = "Option::is_none")]
268    pub background: Option<ImageResponseBackground>,
269    #[serde(skip_serializing_if = "Option::is_none")]
270    pub data: Option<Vec<Image>>,
271    #[serde(skip_serializing_if = "Option::is_none")]
272    pub output_format: Option<ImageOutputFormat>,
273    #[serde(skip_serializing_if = "Option::is_none")]
274    pub quality: Option<ImageResponseQuality>,
275    #[serde(skip_serializing_if = "Option::is_none")]
276    pub size: Option<ImageResponseSize>,
277    #[serde(skip_serializing_if = "Option::is_none")]
278    pub usage: Option<ImageUsage>,
279    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
280    pub extra: Extra,
281}
282
283#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
284pub struct Image {
285    #[serde(skip_serializing_if = "Option::is_none")]
286    pub b64_json: Option<String>,
287    #[serde(skip_serializing_if = "Option::is_none")]
288    pub revised_prompt: Option<String>,
289    #[serde(skip_serializing_if = "Option::is_none")]
290    pub url: Option<String>,
291    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
292    pub extra: Extra,
293}
294
295#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
296pub struct ImageUsage {
297    pub input_tokens: u32,
298    pub input_tokens_details: ImageTokenDetails,
299    pub output_tokens: u32,
300    pub total_tokens: u32,
301    #[serde(skip_serializing_if = "Option::is_none")]
302    pub output_tokens_details: Option<ImageTokenDetails>,
303    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
304    pub extra: Extra,
305}
306
307#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
308pub struct ImageTokenDetails {
309    pub image_tokens: u32,
310    pub text_tokens: u32,
311    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
312    pub extra: Extra,
313}
314
315#[derive(Debug, Clone, PartialEq, Serialize)]
316pub enum ImageStreamEvent {
317    Known(KnownImageStreamEvent),
318    Unknown(UnknownImageStreamEvent),
319}
320
321impl<'de> Deserialize<'de> for ImageStreamEvent {
322    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
323    where
324        D: serde::Deserializer<'de>,
325    {
326        let value = Value::deserialize(deserializer)?;
327        match image_stream_event_type::<D::Error>(&value)? {
328            Some(ImageStreamEventType::Known(_)) => serde_json::from_value(value)
329                .map(Self::Known)
330                .map_err(de::Error::custom),
331            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
332                .map(Self::Unknown)
333                .map_err(de::Error::custom),
334        }
335    }
336}
337
338#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
339#[serde(tag = "type")]
340pub enum KnownImageStreamEvent {
341    #[serde(rename = "image_generation.partial_image")]
342    ImageGenerationPartialImage {
343        b64_json: String,
344        partial_image_index: u32,
345        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
346        extra: Extra,
347    },
348    #[serde(rename = "image_generation.completed")]
349    ImageGenerationCompleted {
350        b64_json: String,
351        #[serde(skip_serializing_if = "Option::is_none")]
352        usage: Option<ImageUsage>,
353        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
354        extra: Extra,
355    },
356    #[serde(rename = "image_edit.partial_image")]
357    ImageEditPartialImage {
358        b64_json: String,
359        partial_image_index: u32,
360        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
361        extra: Extra,
362    },
363    #[serde(rename = "image_edit.completed")]
364    ImageEditCompleted {
365        b64_json: String,
366        #[serde(skip_serializing_if = "Option::is_none")]
367        usage: Option<ImageUsage>,
368        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
369        extra: Extra,
370    },
371}
372
373#[derive(Debug, Clone, PartialEq, Serialize)]
374pub enum ImageGenerationStreamEvent {
375    Known(KnownImageGenerationStreamEvent),
376    Unknown(UnknownImageStreamEvent),
377}
378
379impl<'de> Deserialize<'de> for ImageGenerationStreamEvent {
380    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
381    where
382        D: serde::Deserializer<'de>,
383    {
384        let value = Value::deserialize(deserializer)?;
385        match image_stream_event_type::<D::Error>(&value)? {
386            Some(ImageStreamEventType::Known(
387                ImageStreamEventTypeKnown::ImageGenerationPartialImage
388                | ImageStreamEventTypeKnown::ImageGenerationCompleted,
389            )) => serde_json::from_value(value)
390                .map(Self::Known)
391                .map_err(de::Error::custom),
392            Some(ImageStreamEventType::Known(_)) => Err(de::Error::custom(
393                "known image edit stream event cannot deserialize as image generation stream event",
394            )),
395            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
396                .map(Self::Unknown)
397                .map_err(de::Error::custom),
398        }
399    }
400}
401
402#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
403#[serde(tag = "type")]
404pub enum KnownImageGenerationStreamEvent {
405    #[serde(rename = "image_generation.partial_image")]
406    PartialImage {
407        b64_json: String,
408        partial_image_index: u32,
409        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
410        extra: Extra,
411    },
412    #[serde(rename = "image_generation.completed")]
413    Completed {
414        b64_json: String,
415        #[serde(skip_serializing_if = "Option::is_none")]
416        usage: Option<ImageUsage>,
417        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
418        extra: Extra,
419    },
420}
421
422#[derive(Debug, Clone, PartialEq, Serialize)]
423pub enum ImageEditStreamEvent {
424    Known(KnownImageEditStreamEvent),
425    Unknown(UnknownImageStreamEvent),
426}
427
428impl<'de> Deserialize<'de> for ImageEditStreamEvent {
429    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
430    where
431        D: serde::Deserializer<'de>,
432    {
433        let value = Value::deserialize(deserializer)?;
434        match image_stream_event_type::<D::Error>(&value)? {
435            Some(ImageStreamEventType::Known(
436                ImageStreamEventTypeKnown::ImageEditPartialImage
437                | ImageStreamEventTypeKnown::ImageEditCompleted,
438            )) => serde_json::from_value(value)
439                .map(Self::Known)
440                .map_err(de::Error::custom),
441            Some(ImageStreamEventType::Known(_)) => Err(de::Error::custom(
442                "known image generation stream event cannot deserialize as image edit stream event",
443            )),
444            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
445                .map(Self::Unknown)
446                .map_err(de::Error::custom),
447        }
448    }
449}
450
451#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
452#[serde(tag = "type")]
453pub enum KnownImageEditStreamEvent {
454    #[serde(rename = "image_edit.partial_image")]
455    PartialImage {
456        b64_json: String,
457        partial_image_index: u32,
458        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
459        extra: Extra,
460    },
461    #[serde(rename = "image_edit.completed")]
462    Completed {
463        b64_json: String,
464        #[serde(skip_serializing_if = "Option::is_none")]
465        usage: Option<ImageUsage>,
466        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
467        extra: Extra,
468    },
469}
470
471#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
472pub struct UnknownImageStreamEvent {
473    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
474    pub type_: Option<ImageStreamEventType>,
475    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
476    pub extra: Extra,
477}
478
479fn image_stream_event_type<E>(value: &Value) -> Result<Option<ImageStreamEventType>, E>
480where
481    E: de::Error,
482{
483    let Some(type_name) = value.get("type").and_then(Value::as_str) else {
484        return Ok(None);
485    };
486
487    serde_json::from_value(Value::String(type_name.to_owned()))
488        .map(Some)
489        .map_err(de::Error::custom)
490}
491
492#[cfg(test)]
493mod tests {
494    use super::*;
495
496    #[test]
497    fn image_edit_accepts_generic_multipart_json_shape() {
498        let req: ImageEditRequest = serde_json::from_str(
499            r#"{
500                "image": [
501                    "data:image/png;base64,AAAA",
502                    "file_123"
503                ],
504                "mask": "data:image/png;base64,BBBB",
505                "prompt": "make it blue",
506                "model": "gpt-image-1.5",
507                "n": "2",
508                "stream": "true"
509            }"#,
510        )
511        .unwrap();
512
513        assert_eq!(req.images.len(), 2);
514        assert_eq!(
515            req.images[0].image_url.as_deref(),
516            Some("data:image/png;base64,AAAA")
517        );
518        assert_eq!(req.images[1].file_id.as_deref(), Some("file_123"));
519        assert_eq!(
520            req.mask.as_ref().and_then(|mask| mask.image_url.as_deref()),
521            Some("data:image/png;base64,BBBB")
522        );
523        assert_eq!(req.n, Some(2));
524        assert_eq!(req.stream, Some(true));
525    }
526}
gproxy_protocol/protocol/openai/images.rs

gproxy_protocol/protocol/openai/
images.rs
