rig_core/providers/xai/
image_generation.rs

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

rig_core/providers/xai/
image_generation.rs
