agentai_image/provider/
mod.rs

1//! # Image providers
2//!
3//! [`ImageProvider`] is the interface between [`ImageGenerator`](crate::ImageGenerator) and an
4//! image generation API. Each submodule implements it for one API:
5//!
6//! - [`openai`]: the OpenAI image API and APIs compatible with it.
7//! - [`openrouter`]: the OpenRouter image API, with models of many vendors.
8//!
9//! [`list_models`] lists the image models a provider offers, with names ready for
10//! [`ImageGenerator::new`](crate::ImageGenerator::new).
11//!
12//! Usually you do not pick a provider yourself: [`from_model`] (used by
13//! [`ImageGenerator::new`](crate::ImageGenerator::new)) chooses it from the model name. Build one
14//! yourself to set the API key or base URL in code, or implement [`ImageProvider`] for another
15//! API.
16
17pub mod openai;
18pub mod openrouter;
19
20use agentai_core::HttpError;
21use agentai_core::message::Media;
22use agentai_core::storage::StorageError;
23use async_trait::async_trait;
24use thiserror::Error;
25
26/// A request to generate an image.
27#[derive(Debug, Clone, PartialEq)]
28#[non_exhaustive]
29pub struct ImageRequest {
30    /// What to draw.
31    pub prompt: String,
32    /// Size, in the format of the provider (for example `1024x1024`). Not sent when `None`.
33    pub size: Option<String>,
34    /// Quality, in the format of the provider (for example `high`). Not sent when `None`.
35    pub quality: Option<String>,
36    /// File name of the image, set as [`Media::name`] and so passed to the storage. Not sent
37    /// to the provider. Without an extension, the one of the image type is added.
38    pub name: Option<String>,
39}
40
41impl ImageRequest {
42    /// Creates a request with a prompt.
43    pub fn new(prompt: impl Into<String>) -> Self {
44        Self {
45            prompt: prompt.into(),
46            size: None,
47            quality: None,
48            name: None,
49        }
50    }
51
52    /// Sets the size.
53    pub fn with_size(mut self, size: impl Into<String>) -> Self {
54        self.size = Some(size.into());
55        self
56    }
57
58    /// Sets the quality.
59    pub fn with_quality(mut self, quality: impl Into<String>) -> Self {
60        self.quality = Some(quality.into());
61        self
62    }
63
64    /// Sets the file name of the image, see [`ImageRequest::name`].
65    pub fn with_name(mut self, name: impl Into<String>) -> Self {
66        self.name = Some(name.into());
67        self
68    }
69}
70
71/// A request to edit images, or to make a new image from them.
72#[derive(Debug, Clone, PartialEq)]
73#[non_exhaustive]
74pub struct ImageEditRequest {
75    /// What to change.
76    pub prompt: String,
77    /// The images to edit.
78    pub images: Vec<Media>,
79    /// Size, see [`ImageRequest::size`].
80    pub size: Option<String>,
81    /// Quality, see [`ImageRequest::quality`].
82    pub quality: Option<String>,
83    /// File name of the new image, see [`ImageRequest::name`].
84    pub name: Option<String>,
85}
86
87impl ImageEditRequest {
88    /// Creates a request with a prompt and the images to edit.
89    pub fn new(prompt: impl Into<String>, images: Vec<Media>) -> Self {
90        Self {
91            prompt: prompt.into(),
92            images,
93            size: None,
94            quality: None,
95            name: None,
96        }
97    }
98
99    /// Sets the size.
100    pub fn with_size(mut self, size: impl Into<String>) -> Self {
101        self.size = Some(size.into());
102        self
103    }
104
105    /// Sets the quality.
106    pub fn with_quality(mut self, quality: impl Into<String>) -> Self {
107        self.quality = Some(quality.into());
108        self
109    }
110
111    /// Sets the file name of the image, see [`ImageRequest::name`].
112    pub fn with_name(mut self, name: impl Into<String>) -> Self {
113        self.name = Some(name.into());
114        self
115    }
116}
117
118/// An image made by a provider.
119#[derive(Debug, Clone, PartialEq)]
120#[non_exhaustive]
121pub struct ImageResponse {
122    /// The image, as bytes (base64) or as a URL, depending on the provider.
123    pub image: Media,
124    /// The prompt the provider actually used, if it rewrote yours.
125    pub revised_prompt: Option<String>,
126    /// ID of the image in the storage of the [`ImageGenerator`](crate::ImageGenerator), if
127    /// it has one.
128    pub id: Option<String>,
129}
130
131impl ImageResponse {
132    /// Creates a response with an image.
133    pub fn new(image: Media) -> Self {
134        Self {
135            image,
136            revised_prompt: None,
137            id: None,
138        }
139    }
140
141    /// Sets the revised prompt.
142    pub fn with_revised_prompt(mut self, prompt: impl Into<String>) -> Self {
143        self.revised_prompt = Some(prompt.into());
144        self
145    }
146}
147
148/// An image model offered by a provider, see [`list_models`].
149#[derive(Debug, Clone, PartialEq)]
150#[non_exhaustive]
151pub struct ImageModel {
152    /// Name for [`ImageGenerator::new`](crate::ImageGenerator::new) and [`from_model`], with
153    /// the provider prefix, for example `openrouter::bytedance-seed/seedream-4.5`.
154    pub model: String,
155    /// Name of the model in the API.
156    pub id: String,
157    /// Display name.
158    pub name: Option<String>,
159    /// Description.
160    pub description: Option<String>,
161    /// Whether the model takes input images (needed to edit), when the API says so.
162    pub accepts_images: Option<bool>,
163}
164
165impl ImageModel {
166    /// Creates a model with its full name and its name in the API.
167    pub fn new(model: impl Into<String>, id: impl Into<String>) -> Self {
168        Self {
169            model: model.into(),
170            id: id.into(),
171            name: None,
172            description: None,
173            accepts_images: None,
174        }
175    }
176}
177
178/// Errors of image providers.
179#[derive(Debug, Error)]
180#[non_exhaustive]
181pub enum ImageError {
182    /// The request could not be sent or the response could not be read.
183    #[error(transparent)]
184    Http(#[from] HttpError),
185    /// The API answered with an error status.
186    #[error("image API error (status {status}): {message}")]
187    Api { status: u16, message: String },
188    /// The API answered with something the provider cannot read.
189    #[error("invalid image API response: {0}")]
190    InvalidResponse(String),
191    /// The provider cannot do this, for example edit images.
192    #[error("not supported by the image provider: {0}")]
193    Unsupported(&'static str),
194    /// No provider matches the model name.
195    #[error("no image provider for model '{0}'")]
196    UnknownProvider(String),
197    /// The environment variable with the API key is not set.
198    #[error("the API key is missing, set {0}")]
199    MissingApiKey(&'static str),
200    /// The image was made but could not be saved in the storage.
201    #[error("cannot save the image: {0}")]
202    Storage(#[from] StorageError),
203}
204
205impl ImageError {
206    /// Returns `true` when the API rejected the credentials (status 401 or 403).
207    pub fn is_auth(&self) -> bool {
208        matches!(
209            self,
210            Self::Api {
211                status: 401 | 403,
212                ..
213            }
214        )
215    }
216}
217
218/// An image generation API.
219///
220/// Implement [`generate`](ImageProvider::generate); a provider that can edit images also
221/// overrides [`supports_edit`](ImageProvider::supports_edit) and
222/// [`edit`](ImageProvider::edit).
223#[async_trait]
224pub trait ImageProvider: Send + Sync + 'static {
225    /// Name of the model, if the provider has one.
226    fn model(&self) -> Option<&str> {
227        None
228    }
229
230    /// Generates an image.
231    async fn generate(&self, request: ImageRequest) -> Result<ImageResponse, ImageError>;
232
233    /// Returns `true` when the provider implements [`edit`](ImageProvider::edit).
234    fn supports_edit(&self) -> bool {
235        false
236    }
237
238    /// Edits images. The default implementation returns [`ImageError::Unsupported`].
239    async fn edit(&self, request: ImageEditRequest) -> Result<ImageResponse, ImageError> {
240        let _ = request;
241        Err(ImageError::Unsupported("image editing"))
242    }
243
244    /// Lists the image models of the API. The default implementation returns
245    /// [`ImageError::Unsupported`].
246    async fn list_models(&self) -> Result<Vec<ImageModel>, ImageError> {
247        Err(ImageError::Unsupported("listing models"))
248    }
249}
250
251/// The APIs [`from_model`] knows.
252#[derive(Debug, Clone, Copy, PartialEq, Eq)]
253enum ProviderKind {
254    OpenAi,
255    OpenRouter,
256}
257
258impl ProviderKind {
259    fn from_name(name: &str) -> Option<Self> {
260        match name {
261            "openai" => Some(Self::OpenAi),
262            "openrouter" => Some(Self::OpenRouter),
263            _ => None,
264        }
265    }
266
267    fn build(self, model: &str) -> Result<Box<dyn ImageProvider>, ImageError> {
268        Ok(match self {
269            Self::OpenAi => Box::new(openai::OpenAiImageProvider::from_env(model)?),
270            Self::OpenRouter => Box::new(openrouter::OpenRouterImageProvider::from_env(model)?),
271        })
272    }
273}
274
275/// Splits `provider::model` and finds the provider, like model names in `agentai`.
276fn resolve(model: &str) -> Result<(ProviderKind, &str), ImageError> {
277    let unknown = || ImageError::UnknownProvider(model.to_string());
278    match model.split_once("::") {
279        Some((provider, name)) => ProviderKind::from_name(provider)
280            .map(|kind| (kind, name))
281            .ok_or_else(unknown),
282        None if model.starts_with("gpt-image") || model.starts_with("dall-e") => {
283            Ok((ProviderKind::OpenAi, model))
284        }
285        None => Err(unknown()),
286    }
287}
288
289/// Creates the provider for a model name, with its settings from environment variables.
290///
291/// The provider comes from a `provider::` prefix (`openai::`, `openrouter::`), for example
292/// `openai::gpt-image-1` or `openrouter::bytedance-seed/seedream-4.5`, or, for
293/// well-known names (`gpt-image-*`, `dall-e-*`), from the name alone. Use the prefix for other
294/// models behind an OpenAI-compatible API. See each provider's `from_env` for its variables,
295/// for example [`OpenAiImageProvider::from_env`](openai::OpenAiImageProvider::from_env).
296pub fn from_model(model: &str) -> Result<Box<dyn ImageProvider>, ImageError> {
297    let (kind, name) = resolve(model)?;
298    kind.build(name)
299}
300
301/// Lists the image models of a provider (`openai` or `openrouter`), with its settings from
302/// environment variables as in [`from_model`].
303///
304/// ```rust,no_run
305/// # async fn run() -> Result<(), agentai_image::ImageError> {
306/// for model in agentai_image::provider::list_models("openrouter").await? {
307///     println!("{} {:?}", model.model, model.name);
308/// }
309/// # Ok(())
310/// # }
311/// ```
312pub async fn list_models(provider: &str) -> Result<Vec<ImageModel>, ImageError> {
313    let kind = ProviderKind::from_name(provider)
314        .ok_or_else(|| ImageError::UnknownProvider(provider.to_string()))?;
315    kind.build("")?.list_models().await
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    #[test]
323    fn model_names_select_the_provider() {
324        assert_eq!(
325            resolve("openai::my-model").unwrap(),
326            (ProviderKind::OpenAi, "my-model")
327        );
328        assert_eq!(
329            resolve("gpt-image-1").unwrap(),
330            (ProviderKind::OpenAi, "gpt-image-1")
331        );
332        assert_eq!(
333            resolve("openrouter::black-forest-labs/flux.2").unwrap(),
334            (ProviderKind::OpenRouter, "black-forest-labs/flux.2")
335        );
336        assert_eq!(
337            resolve("dall-e-3").unwrap(),
338            (ProviderKind::OpenAi, "dall-e-3")
339        );
340        for model in ["flux-1", "other::gpt-image-1"] {
341            assert!(matches!(resolve(model), Err(ImageError::UnknownProvider(m)) if m == model));
342        }
343    }
344}
agentai_image/provider/mod.rs

agentai_image/provider/
mod.rs
