agentai_image/provider/
openrouter.rs

1//! # OpenRouter image API
2//!
3//! [`OpenRouterImageProvider`] generates images with `POST /images`, the dedicated image API of
4//! [OpenRouter](https://openrouter.ai), and lists image models with
5//! `GET /models?output_modalities=image`. Images to edit are sent as `input_references`, so
6//! editing works with every model that takes input images.
7
8use super::openai::{parse_response, send_request};
9use super::{ImageEditRequest, ImageError, ImageModel, ImageProvider, ImageRequest, ImageResponse};
10use agentai_core::message::{Media, MediaSource};
11use agentai_core::{HttpClient, HttpRequest, ReqwestClient};
12use async_trait::async_trait;
13use serde::{Deserialize, Serialize};
14use std::sync::Arc;
15
16/// The OpenRouter API endpoint.
17pub const DEFAULT_BASE_URL: &str = "https://openrouter.ai/api/v1";
18/// Environment variable with the API key.
19pub const API_KEY_ENV: &str = "OPENROUTER_API_KEY";
20/// Environment variable with the base URL.
21pub const BASE_URL_ENV: &str = "OPENROUTER_BASE_URL";
22
23/// Provider for the OpenRouter image API.
24///
25/// ```rust,no_run
26/// # use agentai_image::provider::openrouter::{DEFAULT_BASE_URL, OpenRouterImageProvider};
27/// # fn main() -> Result<(), agentai_image::ImageError> {
28/// // From `OPENROUTER_API_KEY` and, if set, `OPENROUTER_BASE_URL`
29/// let provider = OpenRouterImageProvider::from_env("bytedance-seed/seedream-4.5")?;
30/// // Or everything in code
31/// let provider =
32///     OpenRouterImageProvider::new("bytedance-seed/seedream-4.5", "sk-or-...", DEFAULT_BASE_URL);
33/// # Ok(())
34/// # }
35/// ```
36#[derive(Clone)]
37pub struct OpenRouterImageProvider {
38    model: String,
39    api_key: String,
40    base_url: String,
41    client: Arc<dyn HttpClient>,
42}
43
44impl std::fmt::Debug for OpenRouterImageProvider {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("OpenRouterImageProvider")
47            .field("model", &self.model)
48            .field("base_url", &self.base_url)
49            .finish_non_exhaustive()
50    }
51}
52
53impl OpenRouterImageProvider {
54    /// Creates a provider for `model` with the API key and base URL given in code, for example
55    /// [`DEFAULT_BASE_URL`].
56    pub fn new(
57        model: impl Into<String>,
58        api_key: impl Into<String>,
59        base_url: impl Into<String>,
60    ) -> Self {
61        Self {
62            model: model.into(),
63            api_key: api_key.into(),
64            base_url: base_url.into().trim_end_matches('/').to_string(),
65            client: Arc::new(ReqwestClient::new()),
66        }
67    }
68
69    /// Creates a provider for `model` with the API key from `OPENROUTER_API_KEY` and the base
70    /// URL from `OPENROUTER_BASE_URL`, or [`DEFAULT_BASE_URL`] when it is not set.
71    pub fn from_env(model: impl Into<String>) -> Result<Self, ImageError> {
72        let api_key =
73            std::env::var(API_KEY_ENV).map_err(|_| ImageError::MissingApiKey(API_KEY_ENV))?;
74        let base_url = std::env::var(BASE_URL_ENV).unwrap_or_else(|_| DEFAULT_BASE_URL.to_string());
75        Ok(Self::new(model, api_key, base_url))
76    }
77
78    /// Model name sent to the API.
79    pub fn model(&self) -> &str {
80        &self.model
81    }
82
83    #[cfg(test)]
84    fn with_client(mut self, client: Arc<dyn HttpClient>) -> Self {
85        self.client = client;
86        self
87    }
88
89    async fn send(&self, body: &ImageBody<'_>) -> Result<ImageResponse, ImageError> {
90        let request = HttpRequest::post(format!("{}/images", self.base_url)).with_json(body)?;
91        parse_response(&send_request(&*self.client, &self.api_key, request).await?)
92    }
93
94    /// Origin of the base URL, to resolve the relative `links.next` of a model list.
95    fn origin(&self) -> &str {
96        let after_scheme = self.base_url.find("://").map_or(0, |i| i + 3);
97        match self.base_url[after_scheme..].find('/') {
98            Some(i) => &self.base_url[..after_scheme + i],
99            None => &self.base_url,
100        }
101    }
102}
103
104#[derive(Serialize)]
105struct ImageBody<'a> {
106    model: &'a str,
107    prompt: &'a str,
108    #[serde(skip_serializing_if = "Option::is_none")]
109    size: Option<&'a str>,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    quality: Option<&'a str>,
112    #[serde(skip_serializing_if = "Vec::is_empty")]
113    input_references: Vec<ImageReference>,
114}
115
116#[derive(Serialize)]
117struct ImageReference {
118    r#type: &'static str,
119    image_url: ImageUrl,
120}
121
122#[derive(Serialize)]
123struct ImageUrl {
124    url: String,
125}
126
127impl ImageReference {
128    /// A data URL for bytes, the URL itself otherwise.
129    fn new(media: &Media) -> Self {
130        let url = match &media.source {
131            MediaSource::Base64(data) => format!("data:{};base64,{data}", media.content_type),
132            MediaSource::Url(url) => url.clone(),
133        };
134        Self {
135            r#type: "image_url",
136            image_url: ImageUrl { url },
137        }
138    }
139}
140
141#[derive(Deserialize)]
142struct ModelList {
143    #[serde(default)]
144    data: Vec<ModelEntry>,
145    links: Option<Links>,
146}
147
148#[derive(Deserialize)]
149struct Links {
150    next: Option<String>,
151}
152
153#[derive(Deserialize)]
154struct ModelEntry {
155    id: String,
156    name: Option<String>,
157    description: Option<String>,
158    architecture: Option<Architecture>,
159}
160
161#[derive(Deserialize)]
162struct Architecture {
163    input_modalities: Option<Vec<String>>,
164}
165
166#[async_trait]
167impl ImageProvider for OpenRouterImageProvider {
168    fn model(&self) -> Option<&str> {
169        Some(&self.model)
170    }
171
172    async fn generate(&self, request: ImageRequest) -> Result<ImageResponse, ImageError> {
173        self.send(&ImageBody {
174            model: &self.model,
175            prompt: &request.prompt,
176            size: request.size.as_deref(),
177            quality: request.quality.as_deref(),
178            input_references: Vec::new(),
179        })
180        .await
181    }
182
183    fn supports_edit(&self) -> bool {
184        true
185    }
186
187    /// Sends the images as `input_references`; models that do not take input images answer
188    /// with an API error.
189    async fn edit(&self, request: ImageEditRequest) -> Result<ImageResponse, ImageError> {
190        self.send(&ImageBody {
191            model: &self.model,
192            prompt: &request.prompt,
193            size: request.size.as_deref(),
194            quality: request.quality.as_deref(),
195            input_references: request.images.iter().map(ImageReference::new).collect(),
196        })
197        .await
198    }
199
200    /// Lists the models that output images, following every page of the list.
201    async fn list_models(&self) -> Result<Vec<ImageModel>, ImageError> {
202        let mut url = format!("{}/models?output_modalities=image", self.base_url);
203        let mut models = Vec::new();
204        loop {
205            let response =
206                send_request(&*self.client, &self.api_key, HttpRequest::get(&url)).await?;
207            let list: ModelList = response
208                .json()
209                .map_err(|e| ImageError::InvalidResponse(e.to_string()))?;
210            models.extend(list.data.into_iter().map(|entry| {
211                let mut model = ImageModel::new(format!("openrouter::{}", entry.id), entry.id);
212                model.name = entry.name;
213                model.description = entry.description;
214                model.accepts_images = entry
215                    .architecture
216                    .and_then(|a| a.input_modalities)
217                    .map(|inputs| inputs.iter().any(|input| input == "image"));
218                model
219            }));
220            match list.links.and_then(|links| links.next) {
221                Some(next) if next.starts_with('/') => url = format!("{}{next}", self.origin()),
222                Some(next) => url = next,
223                None => return Ok(models),
224            }
225        }
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::provider::openai::tests::{MockHttp, json_response};
233    use agentai_core::HttpRequest;
234    use serde_json::{Value, json};
235
236    fn provider(http: &Arc<MockHttp>) -> OpenRouterImageProvider {
237        OpenRouterImageProvider::new("vendor/model", "key", "https://router.test/api/v1/")
238            .with_client(http.clone())
239    }
240
241    fn body_json(request: &HttpRequest) -> Value {
242        serde_json::from_slice(request.body.as_ref().unwrap()).unwrap()
243    }
244
245    #[tokio::test]
246    async fn generate_posts_to_images_and_reads_the_media_type() {
247        let http = MockHttp::new(vec![json_response(
248            200,
249            json!({"created": 1, "data": [{"b64_json": "aGk=", "media_type": "image/webp"}]}),
250        )]);
251
252        let response = provider(&http)
253            .generate(ImageRequest::new("a cat").with_size("1K"))
254            .await
255            .unwrap();
256
257        let request = http.request(0);
258        assert_eq!(request.url, "https://router.test/api/v1/images");
259        assert_eq!(request.header("Authorization"), Some("Bearer key"));
260        assert_eq!(
261            body_json(&request),
262            json!({"model": "vendor/model", "prompt": "a cat", "size": "1K"})
263        );
264        assert_eq!(response.image, Media::from_base64("image/webp", "aGk="));
265    }
266
267    #[tokio::test]
268    async fn edit_sends_bytes_as_data_urls_and_urls_as_they_are() {
269        let http = MockHttp::new(vec![json_response(
270            200,
271            json!({"data": [{"b64_json": "aGk="}]}),
272        )]);
273        let images = vec![
274            Media::from_base64("image/png", "aGk="),
275            Media::image_url("https://example.com/cat.jpg"),
276        ];
277
278        provider(&http)
279            .edit(ImageEditRequest::new("add a hat", images))
280            .await
281            .unwrap();
282
283        assert_eq!(
284            body_json(&http.request(0))["input_references"],
285            json!([
286                {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGk="}},
287                {"type": "image_url", "image_url": {"url": "https://example.com/cat.jpg"}},
288            ])
289        );
290    }
291
292    #[tokio::test]
293    async fn listing_asks_for_image_models_and_follows_every_page() {
294        let http = MockHttp::new(vec![
295            json_response(
296                200,
297                json!({
298                    "data": [{
299                        "id": "a/edit",
300                        "name": "Edit",
301                        "architecture": {"input_modalities": ["text", "image"]},
302                    }],
303                    "links": {"next": "/api/v1/models?output_modalities=image&offset=1"},
304                }),
305            ),
306            json_response(
307                200,
308                json!({
309                    "data": [{"id": "b/draw", "architecture": {"input_modalities": ["text"]}}],
310                    "links": {"next": null},
311                }),
312            ),
313        ]);
314
315        let models = provider(&http).list_models().await.unwrap();
316
317        assert_eq!(
318            http.request(0).url,
319            "https://router.test/api/v1/models?output_modalities=image"
320        );
321        assert_eq!(
322            http.request(1).url,
323            "https://router.test/api/v1/models?output_modalities=image&offset=1"
324        );
325        let mut edit = ImageModel::new("openrouter::a/edit", "a/edit");
326        edit.name = Some("Edit".to_string());
327        edit.accepts_images = Some(true);
328        let mut draw = ImageModel::new("openrouter::b/draw", "b/draw");
329        draw.accepts_images = Some(false);
330        assert_eq!(models, [edit, draw]);
331    }
332}
agentai_image/provider/openrouter.rs

agentai_image/provider/
openrouter.rs
