rig_core/
image_generation.rs

1//! Everything related to core image generation abstractions in Rig.
2//! Rig allows calling a number of different providers (that support image generation) using the [ImageGenerationModel] trait.
3use crate::markers::{Missing, Provided};
4use crate::{http_client, provider_response};
5use serde_json::Value;
6use thiserror::Error;
7
8/// Errors returned by image generation models.
9///
10/// Inspect provider failures with [`Self::provider_response_body`],
11/// [`Self::provider_response_json`], and [`Self::provider_response_status`].
12#[derive(Debug, Error)]
13#[non_exhaustive]
14pub enum ImageGenerationError {
15    /// Http error (e.g.: connection error, timeout, etc.)
16    #[error("HttpError: {0}")]
17    HttpError(#[from] http_client::Error),
18
19    /// Json error (e.g.: serialization, deserialization)
20    #[error("JsonError: {0}")]
21    JsonError(#[from] serde_json::Error),
22
23    /// Error building the image generation request
24    #[error("RequestError: {0}")]
25    RequestError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
26
27    /// Error parsing the image generation response
28    #[error("ResponseError: {0}")]
29    ResponseError(String),
30
31    /// Error returned by the image generation model provider
32    #[error("ProviderError: {0}")]
33    ProviderError(String),
34
35    /// Raw error response preserved from the image generation model provider
36    #[error("ProviderResponseError: {0}")]
37    ProviderResponse(provider_response::ProviderResponseError),
38}
39
40crate::provider_response::impl_provider_response_helpers!(ImageGenerationError);
41
42/// A unified response for a model image generation, returning both the image and the raw response.
43#[derive(Debug)]
44pub struct ImageGenerationResponse<T> {
45    pub image: Vec<u8>,
46    pub response: T,
47}
48
49pub trait ImageGenerationModel: Clone + Send + Sync {
50    type Response: Send + Sync;
51
52    type Client;
53
54    fn make(client: &Self::Client, model: impl Into<String>) -> Self;
55
56    fn image_generation(
57        &self,
58        request: ImageGenerationRequest,
59    ) -> impl std::future::Future<
60        Output = Result<ImageGenerationResponse<Self::Response>, ImageGenerationError>,
61    > + Send;
62
63    fn image_generation_request(&self) -> ImageGenerationRequestBuilder<Self, Missing> {
64        ImageGenerationRequestBuilder::new(self.clone())
65    }
66}
67/// An image generation request.
68#[non_exhaustive]
69pub struct ImageGenerationRequest {
70    pub prompt: String,
71    pub width: u32,
72    pub height: u32,
73    pub additional_params: Option<Value>,
74}
75
76/// A builder for `ImageGenerationRequest`.
77/// Can be sent to a model provider.
78#[non_exhaustive]
79pub struct ImageGenerationRequestBuilder<M, P = Missing>
80where
81    M: ImageGenerationModel,
82{
83    model: M,
84    prompt: P,
85    width: u32,
86    height: u32,
87    additional_params: Option<Value>,
88}
89
90impl<M> ImageGenerationRequestBuilder<M, Missing>
91where
92    M: ImageGenerationModel,
93{
94    pub fn new(model: M) -> Self {
95        Self {
96            model,
97            prompt: Missing,
98            height: 256,
99            width: 256,
100            additional_params: None,
101        }
102    }
103}
104
105impl<M, P> ImageGenerationRequestBuilder<M, P>
106where
107    M: ImageGenerationModel,
108{
109    /// Sets the prompt for the image generation request
110    pub fn prompt(self, prompt: &str) -> ImageGenerationRequestBuilder<M, Provided<String>> {
111        ImageGenerationRequestBuilder {
112            model: self.model,
113            prompt: Provided(prompt.to_string()),
114            width: self.width,
115            height: self.height,
116            additional_params: self.additional_params,
117        }
118    }
119
120    /// The width of the generated image
121    pub fn width(mut self, width: u32) -> Self {
122        self.width = width;
123        self
124    }
125
126    /// The height of the generated image
127    pub fn height(mut self, height: u32) -> Self {
128        self.height = height;
129        self
130    }
131
132    /// Adds additional parameters to the image generation request.
133    pub fn additional_params(mut self, params: Value) -> Self {
134        self.additional_params = Some(params);
135        self
136    }
137}
138
139impl<M> ImageGenerationRequestBuilder<M, Provided<String>>
140where
141    M: ImageGenerationModel,
142{
143    pub fn build(self) -> ImageGenerationRequest {
144        ImageGenerationRequest {
145            prompt: self.prompt.0,
146            width: self.width,
147            height: self.height,
148            additional_params: self.additional_params,
149        }
150    }
151
152    pub async fn send(self) -> Result<ImageGenerationResponse<M::Response>, ImageGenerationError> {
153        let model = self.model.clone();
154
155        model.image_generation(self.build()).await
156    }
157}
158
159#[cfg(test)]
160mod provider_response_tests {
161    use super::*;
162    use http::StatusCode;
163
164    #[test]
165    fn image_generation_error_provider_response_helpers_with_preserved_json_body() {
166        let body = r#"{"error":{"message":"content policy"}}"#;
167        let error =
168            ImageGenerationError::ProviderResponse(provider_response::ProviderResponseError {
169                status: None,
170                body: body.to_string(),
171            });
172
173        assert_eq!(error.provider_response_body(), Some(body));
174        assert_eq!(error.provider_response_status(), None);
175        assert_eq!(
176            error.provider_response_json().expect("valid JSON"),
177            Some(serde_json::json!({ "error": { "message": "content policy" } }))
178        );
179    }
180
181    #[test]
182    fn image_generation_error_provider_response_helpers_with_http_non_success() {
183        let body = r#"{"error":{"message":"bad request"}}"#;
184        let error =
185            ImageGenerationError::HttpError(http_client::Error::InvalidStatusCodeWithMessage(
186                StatusCode::BAD_REQUEST,
187                body.to_string(),
188            ));
189
190        assert_eq!(error.provider_response_body(), Some(body));
191        assert_eq!(
192            error.provider_response_status(),
193            Some(StatusCode::BAD_REQUEST)
194        );
195        assert_eq!(
196            error.provider_response_json().expect("valid JSON"),
197            Some(serde_json::json!({ "error": { "message": "bad request" } }))
198        );
199    }
200
201    #[test]
202    fn image_generation_error_provider_error_is_not_a_provider_response() {
203        let error = ImageGenerationError::ProviderError("internal diagnostic".to_string());
204
205        assert_eq!(error.provider_response_body(), None);
206        assert_eq!(error.provider_response_status(), None);
207        assert_eq!(error.provider_response_json().expect("no body"), None);
208    }
209
210    #[test]
211    fn image_generation_error_provider_response_helpers_with_unrelated_variant() {
212        let error = ImageGenerationError::ResponseError("parse failed".to_string());
213
214        assert_eq!(error.provider_response_body(), None);
215        assert_eq!(error.provider_response_status(), None);
216        assert_eq!(error.provider_response_json().expect("no body"), None);
217    }
218}
rig_core/image_generation.rs

rig_core/
image_generation.rs
