gproxy_protocol/protocol/openai/images/
stream.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize, de};
4use serde_json::Value;
5
6use super::super::common::*;
7use super::responses::ImageUsage;
8
9#[derive(Debug, Clone, PartialEq, Serialize)]
10pub enum ImageStreamEvent {
11    Known(KnownImageStreamEvent),
12    Unknown(UnknownImageStreamEvent),
13}
14
15impl<'de> Deserialize<'de> for ImageStreamEvent {
16    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
17    where
18        D: serde::Deserializer<'de>,
19    {
20        let value = Value::deserialize(deserializer)?;
21        match image_stream_event_type::<D::Error>(&value)? {
22            Some(ImageStreamEventType::Known(_)) => serde_json::from_value(value)
23                .map(Self::Known)
24                .map_err(de::Error::custom),
25            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
26                .map(Self::Unknown)
27                .map_err(de::Error::custom),
28        }
29    }
30}
31
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33#[serde(tag = "type")]
34pub enum KnownImageStreamEvent {
35    #[serde(rename = "image_generation.partial_image")]
36    ImageGenerationPartialImage {
37        b64_json: String,
38        partial_image_index: u32,
39        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
40        extra: Extra,
41    },
42    #[serde(rename = "image_generation.completed")]
43    ImageGenerationCompleted {
44        b64_json: String,
45        #[serde(skip_serializing_if = "Option::is_none")]
46        usage: Option<ImageUsage>,
47        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
48        extra: Extra,
49    },
50    #[serde(rename = "image_edit.partial_image")]
51    ImageEditPartialImage {
52        b64_json: String,
53        partial_image_index: u32,
54        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
55        extra: Extra,
56    },
57    #[serde(rename = "image_edit.completed")]
58    ImageEditCompleted {
59        b64_json: String,
60        #[serde(skip_serializing_if = "Option::is_none")]
61        usage: Option<ImageUsage>,
62        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
63        extra: Extra,
64    },
65}
66
67#[derive(Debug, Clone, PartialEq, Serialize)]
68pub enum ImageGenerationStreamEvent {
69    Known(KnownImageGenerationStreamEvent),
70    Unknown(UnknownImageStreamEvent),
71}
72
73impl<'de> Deserialize<'de> for ImageGenerationStreamEvent {
74    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
75    where
76        D: serde::Deserializer<'de>,
77    {
78        let value = Value::deserialize(deserializer)?;
79        match image_stream_event_type::<D::Error>(&value)? {
80            Some(ImageStreamEventType::Known(
81                ImageStreamEventTypeKnown::ImageGenerationPartialImage
82                | ImageStreamEventTypeKnown::ImageGenerationCompleted,
83            )) => serde_json::from_value(value)
84                .map(Self::Known)
85                .map_err(de::Error::custom),
86            Some(ImageStreamEventType::Known(_)) => Err(de::Error::custom(
87                "known image edit stream event cannot deserialize as image generation stream event",
88            )),
89            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
90                .map(Self::Unknown)
91                .map_err(de::Error::custom),
92        }
93    }
94}
95
96#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97#[serde(tag = "type")]
98pub enum KnownImageGenerationStreamEvent {
99    #[serde(rename = "image_generation.partial_image")]
100    PartialImage {
101        b64_json: String,
102        partial_image_index: u32,
103        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
104        extra: Extra,
105    },
106    #[serde(rename = "image_generation.completed")]
107    Completed {
108        b64_json: String,
109        #[serde(skip_serializing_if = "Option::is_none")]
110        usage: Option<ImageUsage>,
111        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
112        extra: Extra,
113    },
114}
115
116#[derive(Debug, Clone, PartialEq, Serialize)]
117pub enum ImageEditStreamEvent {
118    Known(KnownImageEditStreamEvent),
119    Unknown(UnknownImageStreamEvent),
120}
121
122impl<'de> Deserialize<'de> for ImageEditStreamEvent {
123    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
124    where
125        D: serde::Deserializer<'de>,
126    {
127        let value = Value::deserialize(deserializer)?;
128        match image_stream_event_type::<D::Error>(&value)? {
129            Some(ImageStreamEventType::Known(
130                ImageStreamEventTypeKnown::ImageEditPartialImage
131                | ImageStreamEventTypeKnown::ImageEditCompleted,
132            )) => serde_json::from_value(value)
133                .map(Self::Known)
134                .map_err(de::Error::custom),
135            Some(ImageStreamEventType::Known(_)) => Err(de::Error::custom(
136                "known image generation stream event cannot deserialize as image edit stream event",
137            )),
138            Some(ImageStreamEventType::Unknown(_)) | None => serde_json::from_value(value)
139                .map(Self::Unknown)
140                .map_err(de::Error::custom),
141        }
142    }
143}
144
145#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
146#[serde(tag = "type")]
147pub enum KnownImageEditStreamEvent {
148    #[serde(rename = "image_edit.partial_image")]
149    PartialImage {
150        b64_json: String,
151        partial_image_index: u32,
152        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
153        extra: Extra,
154    },
155    #[serde(rename = "image_edit.completed")]
156    Completed {
157        b64_json: String,
158        #[serde(skip_serializing_if = "Option::is_none")]
159        usage: Option<ImageUsage>,
160        #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
161        extra: Extra,
162    },
163}
164
165#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
166pub struct UnknownImageStreamEvent {
167    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
168    pub type_: Option<ImageStreamEventType>,
169    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
170    pub extra: Extra,
171}
172
173fn image_stream_event_type<E>(value: &Value) -> Result<Option<ImageStreamEventType>, E>
174where
175    E: de::Error,
176{
177    let Some(type_name) = value.get("type").and_then(Value::as_str) else {
178        return Ok(None);
179    };
180
181    serde_json::from_value(Value::String(type_name.to_owned()))
182        .map(Some)
183        .map_err(de::Error::custom)
184}
gproxy_protocol/protocol/openai/images/stream.rs

gproxy_protocol/protocol/openai/images/
stream.rs
