rig_core/providers/openai/
image_generation.rs

1use super::{Client, client::ApiResponse};
2use crate::http_client::HttpClientExt;
3use crate::image_generation::{ImageGenerationError, ImageGenerationRequest};
4use crate::json_utils::merge_inplace;
5use crate::{http_client, image_generation};
6use base64::Engine;
7use base64::prelude::BASE64_STANDARD;
8use serde::Deserialize;
9use serde_json::json;
10
11// ================================================================
12// OpenAI Image Generation API
13// ================================================================
14pub const DALL_E_2: &str = "dall-e-2";
15pub const DALL_E_3: &str = "dall-e-3";
16pub const GPT_IMAGE_1: &str = "gpt-image-1";
17pub const GPT_IMAGE_1_5: &str = "gpt-image-1.5";
18pub const GPT_IMAGE_2: &str = "gpt-image-2";
19
20#[derive(Debug, Deserialize)]
21pub struct ImageGenerationData {
22    pub b64_json: String,
23}
24
25#[derive(Debug, Deserialize)]
26pub struct ImageGenerationResponse {
27    pub created: i32,
28    pub data: Vec<ImageGenerationData>,
29}
30
31impl TryFrom<ImageGenerationResponse>
32    for image_generation::ImageGenerationResponse<ImageGenerationResponse>
33{
34    type Error = ImageGenerationError;
35
36    fn try_from(value: ImageGenerationResponse) -> Result<Self, Self::Error> {
37        let b64_json = value
38            .data
39            .first()
40            .ok_or_else(|| ImageGenerationError::ResponseError("missing image data".into()))?
41            .b64_json
42            .clone();
43
44        let bytes = BASE64_STANDARD
45            .decode(&b64_json)
46            .map_err(|err| ImageGenerationError::ResponseError(err.to_string()))?;
47
48        Ok(image_generation::ImageGenerationResponse {
49            image: bytes,
50            response: value,
51        })
52    }
53}
54
55#[derive(Clone)]
56pub struct ImageGenerationModel<T = reqwest::Client> {
57    client: Client<T>,
58    /// Name of the model (e.g.: dall-e-2)
59    pub model: String,
60}
61
62impl<T> ImageGenerationModel<T> {
63    pub(crate) fn new(client: Client<T>, model: impl Into<String>) -> Self {
64        Self {
65            client,
66            model: model.into(),
67        }
68    }
69}
70
71impl<T> image_generation::ImageGenerationModel for ImageGenerationModel<T>
72where
73    T: HttpClientExt + Clone + Default + std::fmt::Debug + Send + 'static,
74{
75    type Response = ImageGenerationResponse;
76
77    type Client = Client<T>;
78
79    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
80        Self::new(client.clone(), model)
81    }
82
83    async fn image_generation(
84        &self,
85        generation_request: ImageGenerationRequest,
86    ) -> Result<image_generation::ImageGenerationResponse<Self::Response>, ImageGenerationError>
87    {
88        let mut request = json!({
89            "model": self.model,
90            "prompt": generation_request.prompt,
91            "size": format!("{}x{}", generation_request.width, generation_request.height),
92        });
93
94        if !matches!(
95            self.model.as_str(),
96            GPT_IMAGE_1 | GPT_IMAGE_1_5 | GPT_IMAGE_2
97        ) {
98            merge_inplace(
99                &mut request,
100                json!({
101                    "response_format": "b64_json"
102                }),
103            );
104        }
105
106        let body = serde_json::to_vec(&request)?;
107
108        let request = self
109            .client
110            .post("/images/generations")?
111            .body(body)
112            .map_err(|e| ImageGenerationError::HttpError(e.into()))?;
113
114        let response = self.client.send(request).await?;
115
116        let status = response.status();
117        if !status.is_success() {
118            let text = http_client::text(response).await?;
119
120            return Err(ImageGenerationError::from_http_response(status, text));
121        }
122
123        let text = http_client::text(response).await?;
124
125        match serde_json::from_str::<ApiResponse<ImageGenerationResponse>>(&text)? {
126            ApiResponse::Ok(response) => response.try_into(),
127            ApiResponse::Err(err) => {
128                tracing::warn!(message = %err.message, "provider returned an error response");
129                Err(ImageGenerationError::from_http_response(status, text))
130            }
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use crate::client::image_generation::ImageGenerationClient;
139    use crate::image_generation::ImageGenerationModel as _;
140    use crate::test_utils::RecordingHttpClient;
141
142    fn request() -> ImageGenerationRequest {
143        ImageGenerationRequest {
144            prompt: "draw a cat".to_string(),
145            width: 256,
146            height: 256,
147            additional_params: None,
148        }
149    }
150
151    #[tokio::test]
152    async fn image_generation_non_success_response_preserves_status_and_body() {
153        let body = r#"{"error":{"message":"invalid image","type":"invalid_request_error"}}"#;
154        let http_client =
155            RecordingHttpClient::with_error_response(http::StatusCode::BAD_REQUEST, body);
156        let client = Client::builder()
157            .api_key("test-key")
158            .http_client(http_client)
159            .build()
160            .expect("build client");
161        let model = client.image_generation_model(DALL_E_3);
162
163        let error = model
164            .image_generation(request())
165            .await
166            .expect_err("image generation should fail with non-success status");
167
168        assert!(matches!(error, ImageGenerationError::HttpError(_)));
169        assert_eq!(
170            error.provider_response_status(),
171            Some(http::StatusCode::BAD_REQUEST)
172        );
173        assert_eq!(error.provider_response_body(), Some(body));
174    }
175
176    #[tokio::test]
177    async fn image_generation_preserves_raw_provider_error_json_on_api_error_envelope() {
178        let body = r#"{"message":"quota exceeded","type":"insufficient_quota"}"#;
179        let http_client = RecordingHttpClient::new(body);
180        let client = Client::builder()
181            .api_key("test-key")
182            .http_client(http_client)
183            .build()
184            .expect("build client");
185        let model = client.image_generation_model(DALL_E_3);
186
187        let error = model
188            .image_generation(request())
189            .await
190            .expect_err("image generation should fail with provider error envelope");
191
192        match &error {
193            ImageGenerationError::ProviderResponse(stored) => {
194                assert_eq!(stored.body, body);
195                assert_eq!(stored.status, Some(http::StatusCode::OK));
196                assert_eq!(error.provider_response_body(), Some(body));
197            }
198            other => panic!("expected ProviderResponse, got {other:?}"),
199        }
200    }
201}
rig_core/providers/openai/image_generation.rs

rig_core/providers/openai/
image_generation.rs
