agentai_image/
toolbox.rs

1//! Agent tools (feature `toolbox`): [`GenerateImageToolBox`] and [`ReadMediaToolBox`].
2
3use crate::generator::ImageGenerator;
4use crate::provider::{ImageEditRequest, ImageError, ImageRequest, ImageResponse};
5use agentai::tool::{Tool, ToolBox, ToolError, ToolOutput, ToolResult, toolbox};
6use agentai_core::message::MediaSource;
7use agentai_core::storage::{MediaRead, MediaStorage, StorageError};
8use serde::Serialize;
9use std::sync::Arc;
10
11/// What [`GenerateImageToolBox`] sends back to the model.
12#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13#[non_exhaustive]
14pub enum ImageToolOutput {
15    /// The image itself and a short text with its ID. The model must accept images.
16    #[default]
17    Media,
18    /// Only a JSON text with what was made (ID, type, size in bytes, model, prompts, ...),
19    /// without the image. Saves tokens, and the model can still look at the image with
20    /// [`ReadMediaToolBox`].
21    Metadata,
22}
23
24/// [`ToolBox`] that lets an agent generate and edit images with an [`ImageGenerator`].
25///
26/// Tools:
27/// - `generate_image`: draws an image from a prompt.
28/// - `edit_image`: changes stored images the model names by ID. It needs a storage
29///   ([`with_storage`](GenerateImageToolBox::with_storage)) and a provider that can edit;
30///   without them the model gets an error.
31///
32/// Both take a `name` from the model (for example `lighthouse`), set as
33/// [`Media::name`](agentai_core::message::Media::name) and so passed to
34/// [`MediaSave`](agentai_core::storage::MediaSave); the storage chooses the ID and may use it.
35///
36/// With a storage every image is saved and the model gets its ID. What else the model gets
37/// is set with [`with_output`](GenerateImageToolBox::with_output).
38///
39/// ```rust,no_run
40/// # use agentai_image::{GenerateImageToolBox, ImageGenerator};
41/// # use agentai_image::storage::MemoryStorage;
42/// # fn main() -> Result<(), agentai_image::ImageError> {
43/// let storage = MemoryStorage::new();
44/// let images = GenerateImageToolBox::new(ImageGenerator::new("openai::gpt-image-1")?)
45///     .with_storage(storage.clone());
46/// # Ok(())
47/// # }
48/// ```
49pub struct GenerateImageToolBox {
50    generator: ImageGenerator,
51    storage: Option<Arc<dyn MediaStorage>>,
52    output: ImageToolOutput,
53}
54
55impl GenerateImageToolBox {
56    /// Creates the toolbox. A storage set on `generator` is used to save the images, but
57    /// `edit_image` needs one set with [`with_storage`](GenerateImageToolBox::with_storage).
58    pub fn new(generator: ImageGenerator) -> Self {
59        Self {
60            generator,
61            storage: None,
62            output: ImageToolOutput::default(),
63        }
64    }
65
66    /// Saves every image in `storage` and loads the images to edit from it.
67    pub fn with_storage(mut self, storage: impl MediaStorage) -> Self {
68        let storage: Arc<dyn MediaStorage> = Arc::new(storage);
69        self.generator = self.generator.with_storage(storage.clone());
70        self.storage = Some(storage);
71        self
72    }
73
74    /// Sets what the model gets back, see [`ImageToolOutput`].
75    pub fn with_output(mut self, output: ImageToolOutput) -> Self {
76        self.output = output;
77        self
78    }
79
80    fn output(&self, response: ImageResponse, prompt: &str) -> ToolResult<ToolOutput> {
81        match self.output {
82            ImageToolOutput::Metadata => {
83                let metadata = ImageMetadata {
84                    id: response.id.as_deref(),
85                    content_type: &response.image.content_type,
86                    url: match &response.image.source {
87                        MediaSource::Url(url) => Some(url),
88                        MediaSource::Base64(_) => None,
89                    },
90                    bytes: match &response.image.source {
91                        MediaSource::Base64(data) => Some(base64_len(data)),
92                        MediaSource::Url(_) => None,
93                    },
94                    model: self.generator.model(),
95                    prompt,
96                    revised_prompt: response.revised_prompt.as_deref(),
97                    size: self.generator.size(),
98                    quality: self.generator.quality(),
99                };
100                let json = serde_json::to_string(&metadata).map_err(ToolError::fatal)?;
101                Ok(ToolOutput::from_text(json))
102            }
103            _ => {
104                let mut text = match &response.id {
105                    Some(id) => format!("Created the image with ID {id}."),
106                    None => "Created the image.".to_string(),
107                };
108                if let Some(revised) = &response.revised_prompt {
109                    text.push_str(&format!(" Prompt used: {revised}"));
110                }
111                Ok(ToolOutput::from_text(text).with_media(response.image))
112            }
113        }
114    }
115}
116
117#[toolbox]
118impl GenerateImageToolBox {
119    /// Generates an image from a text prompt.
120    #[tool]
121    async fn generate_image(
122        &self,
123        /// Detailed description of the image: subject, style, composition, colors, text in it.
124        prompt: String,
125        /// Short file name for the new image, for example `lighthouse`. The storage may add to
126        /// it to make the ID unique.
127        name: String,
128    ) -> ToolResult<ToolOutput> {
129        let response = self
130            .generator
131            .generate_with(ImageRequest::new(&prompt).with_name(name))
132            .await
133            .map_err(tool_error)?;
134        self.output(response, &prompt)
135    }
136
137    /// Edits existing images, given by ID, following a text prompt, and creates a new image.
138    #[tool]
139    async fn edit_image(
140        &self,
141        /// What to change, or how to combine the images.
142        prompt: String,
143        /// IDs of the images to edit, as returned by `generate_image` or `edit_image`.
144        images: Vec<String>,
145        /// Short file name for the new image, for example `lighthouse-at-night`. The storage
146        /// may add to it to make the ID unique.
147        name: String,
148    ) -> ToolResult<ToolOutput> {
149        let Some(storage) = &self.storage else {
150            return Err(ToolError::Failed(
151                "editing is not available: there is no storage for the images".to_string(),
152            ));
153        };
154        if images.is_empty() {
155            return Err(ToolError::InvalidArguments(
156                "give at least one image ID in `images`".to_string(),
157            ));
158        }
159        let mut media = Vec::with_capacity(images.len());
160        for id in &images {
161            media.push(storage.read(id).await.map_err(read_error)?);
162        }
163        let response = self
164            .generator
165            .edit_with(ImageEditRequest::new(&prompt, media).with_name(name))
166            .await
167            .map_err(tool_error)?;
168        self.output(response, &prompt)
169    }
170}
171
172/// What [`ImageToolOutput::Metadata`] sends; `None` fields are left out.
173#[derive(Serialize)]
174struct ImageMetadata<'a> {
175    #[serde(skip_serializing_if = "Option::is_none")]
176    id: Option<&'a str>,
177    content_type: &'a str,
178    #[serde(skip_serializing_if = "Option::is_none")]
179    url: Option<&'a String>,
180    #[serde(skip_serializing_if = "Option::is_none")]
181    bytes: Option<usize>,
182    #[serde(skip_serializing_if = "Option::is_none")]
183    model: Option<&'a str>,
184    prompt: &'a str,
185    #[serde(skip_serializing_if = "Option::is_none")]
186    revised_prompt: Option<&'a str>,
187    #[serde(skip_serializing_if = "Option::is_none")]
188    size: Option<&'a str>,
189    #[serde(skip_serializing_if = "Option::is_none")]
190    quality: Option<&'a str>,
191}
192
193/// Size of base64 data once decoded, without decoding it.
194fn base64_len(data: &str) -> usize {
195    let padding = data.bytes().rev().take_while(|&b| b == b'=').count();
196    data.len() / 4 * 3 - padding
197}
198
199/// Rejected credentials cannot be fixed by the model; anything else (a refused prompt, a
200/// provider that cannot edit, a failed request or save) is reported to it.
201fn tool_error(err: ImageError) -> ToolError {
202    if err.is_auth() || matches!(err, ImageError::MissingApiKey(_)) {
203        ToolError::Configuration(err.to_string())
204    } else {
205        ToolError::Failed(err.to_string())
206    }
207}
208
209/// An unknown ID is the model's mistake, so it can correct it.
210fn read_error(err: StorageError) -> ToolError {
211    match err {
212        StorageError::NotFound(id) => {
213            ToolError::InvalidArguments(format!("there is no file with ID '{id}'"))
214        }
215        err => ToolError::Failed(err.to_string()),
216    }
217}
218
219/// [`ToolBox`] with the tool `read_media`, which gives the model a stored file (an image, a
220/// document, ...) by its ID, read from any [`MediaRead`].
221///
222/// Share one storage with [`GenerateImageToolBox`] to let the model look at images it made
223/// with [`ImageToolOutput::Metadata`]:
224///
225/// ```rust
226/// # use agentai_image::ReadMediaToolBox;
227/// # use agentai_image::storage::MemoryStorage;
228/// let storage = MemoryStorage::new();
229/// let files = ReadMediaToolBox::new(storage.clone());
230/// ```
231pub struct ReadMediaToolBox {
232    storage: Arc<dyn MediaRead>,
233}
234
235impl ReadMediaToolBox {
236    /// Creates the toolbox for files in `storage`.
237    pub fn new(storage: impl MediaRead) -> Self {
238        Self {
239            storage: Arc::new(storage),
240        }
241    }
242}
243
244#[toolbox]
245impl ReadMediaToolBox {
246    /// Returns a stored file, for example an image, by its ID, so you can look at it.
247    #[tool]
248    async fn read_media(
249        &self,
250        /// ID of the file.
251        id: String,
252    ) -> ToolResult<ToolOutput> {
253        let media = self.storage.read(&id).await.map_err(read_error)?;
254        let text = format!("File {id} ({})", media.content_type);
255        Ok(ToolOutput::from_text(text).with_media(media))
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262    use crate::provider::openai::tests::{MockHttp, json_response, provider};
263    use agentai_core::message::{Media, Part};
264    use agentai_core::storage::{MediaSave, MemoryStorage};
265    use serde_json::{Value, json};
266
267    fn tool_call(name: &str, arguments: Value) -> agentai::ToolCall {
268        agentai::ToolCall::new("call_1", name, arguments)
269    }
270
271    fn image_reply() -> agentai_core::HttpResponse {
272        json_response(
273            200,
274            json!({"data": [{"b64_json": "aGk=", "revised_prompt": "a red cat"}]}),
275        )
276    }
277
278    fn toolbox(http: &Arc<MockHttp>) -> GenerateImageToolBox {
279        GenerateImageToolBox::new(
280            ImageGenerator::from_provider(provider(http)).with_size("512x512"),
281        )
282    }
283
284    #[tokio::test]
285    async fn generate_image_returns_the_image_with_its_storage_id() {
286        let http = MockHttp::new(vec![image_reply()]);
287        let storage = MemoryStorage::new();
288        let toolbox = toolbox(&http).with_storage(storage.clone());
289
290        let output = toolbox
291            .call_tool(&tool_call(
292                "generate_image",
293                json!({"prompt": "a cat", "name": "image"}),
294            ))
295            .await
296            .unwrap();
297
298        let image = Media::from_base64("image/png", "aGk=").with_name("image.png");
299        assert_eq!(
300            output.parts,
301            vec![
302                Part::Text("Created the image with ID 1-image.png. Prompt used: a red cat".into()),
303                Part::Media(image.clone()),
304            ]
305        );
306        assert_eq!(storage.read("1-image.png").await.unwrap(), image);
307    }
308
309    #[tokio::test]
310    async fn the_name_given_by_the_model_reaches_the_storage() {
311        let http = MockHttp::new(vec![image_reply(), image_reply()]);
312        let storage = MemoryStorage::new();
313        let toolbox = toolbox(&http).with_storage(storage.clone());
314
315        toolbox
316            .call_tool(&tool_call(
317                "generate_image",
318                json!({"prompt": "a cat", "name": "cat"}),
319            ))
320            .await
321            .unwrap();
322        toolbox
323            .call_tool(&tool_call(
324                "edit_image",
325                json!({"prompt": "add a hat", "images": ["1-cat.png"], "name": "cat-with-hat"}),
326            ))
327            .await
328            .unwrap();
329
330        let mut ids = storage.ids();
331        ids.sort();
332        assert_eq!(ids, ["1-cat.png", "2-cat-with-hat.png"]);
333    }
334
335    #[tokio::test]
336    async fn metadata_output_describes_the_image_without_it() {
337        let http = MockHttp::new(vec![image_reply()]);
338        let toolbox = toolbox(&http)
339            .with_storage(MemoryStorage::new())
340            .with_output(ImageToolOutput::Metadata);
341
342        let output = toolbox
343            .call_tool(&tool_call(
344                "generate_image",
345                json!({"prompt": "a cat", "name": "image"}),
346            ))
347            .await
348            .unwrap();
349
350        assert!(output.media().is_empty());
351        let metadata: Value = serde_json::from_str(&output.text().unwrap()).unwrap();
352        assert_eq!(
353            metadata,
354            json!({
355                "id": "1-image.png",
356                "content_type": "image/png",
357                "bytes": 2,
358                "model": "gpt-image-1",
359                "prompt": "a cat",
360                "revised_prompt": "a red cat",
361                "size": "512x512",
362            })
363        );
364    }
365
366    #[tokio::test]
367    async fn edit_image_needs_a_storage_and_loads_the_images_from_it() {
368        let http = MockHttp::new(vec![image_reply()]);
369        let without = toolbox(&http);
370        assert!(matches!(
371            without
372                .call_tool(&tool_call(
373                    "edit_image",
374                    json!({"prompt": "add a hat", "images": ["1"], "name": "hat"})
375                ))
376                .await,
377            Err(ToolError::Failed(_))
378        ));
379
380        let storage = MemoryStorage::new();
381        let id = storage
382            .save(&Media::from_bytes("image/png", b"PNG"))
383            .await
384            .unwrap();
385        let toolbox = toolbox(&http).with_storage(storage);
386        let unknown = toolbox
387            .call_tool(&tool_call(
388                "edit_image",
389                json!({"prompt": "add a hat", "images": ["nope"], "name": "hat"}),
390            ))
391            .await;
392        assert!(matches!(unknown, Err(ToolError::InvalidArguments(_))));
393        let output = toolbox
394            .call_tool(&tool_call(
395                "edit_image",
396                json!({"prompt": "add a hat", "images": [id], "name": "hat"}),
397            ))
398            .await
399            .unwrap();
400
401        assert_eq!(output.media().len(), 1);
402        let request = http.request(0);
403        assert!(request.url.ends_with("/images/edits"));
404        assert!(String::from_utf8_lossy(&request.body.unwrap()).contains("\r\n\r\nPNG\r\n"));
405    }
406
407    #[tokio::test]
408    async fn rejected_keys_stop_the_run_and_other_api_errors_reach_the_model() {
409        let http = MockHttp::new(vec![
410            json_response(401, json!({"error": {"message": "bad key"}})),
411            json_response(400, json!({"error": {"message": "refused"}})),
412        ]);
413        let toolbox = toolbox(&http);
414        let request = tool_call("generate_image", json!({"prompt": "x", "name": "x"}));
415        let call = || toolbox.call_tool(&request);
416
417        assert!(matches!(call().await, Err(ToolError::Configuration(_))));
418        assert!(matches!(call().await, Err(ToolError::Failed(m)) if m.contains("refused")));
419    }
420
421    #[tokio::test]
422    async fn read_media_returns_the_stored_file() {
423        let storage = MemoryStorage::new();
424        let pdf = Media::from_bytes("application/pdf", b"%PDF");
425        let id = storage.save(&pdf).await.unwrap();
426        let toolbox = ReadMediaToolBox::new(storage);
427
428        let output = toolbox
429            .call_tool(&tool_call("read_media", json!({ "id": id })))
430            .await
431            .unwrap();
432        assert_eq!(output.media(), vec![&pdf]);
433        assert!(matches!(
434            toolbox
435                .call_tool(&tool_call("read_media", json!({"id": "nope"})))
436                .await,
437            Err(ToolError::InvalidArguments(_))
438        ));
439    }
440}
agentai_image/toolbox.rs

agentai_image/
toolbox.rs
