rpi_tools/tools/
image.rs

1//! Image processing tool for the agent.
2//!
3//! Provides image processing capabilities: resize, convert, and get info.
4//! Mirrors the image processing utilities from upstream.
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use base64::{engine::general_purpose::STANDARD, Engine};
10use rpi_agent::agent_tool::AgentTool;
11use rpi_agent::error::AgentError;
12use rpi_agent::types::{AgentToolResult, TextContentOrImage, ToolResultPartial};
13use rpi_ai::types::{ConstrainedSamplingConfig, ConstrainedStrictness, Schema, Tool};
14use schemars::JsonSchema;
15use serde::Deserialize;
16use tokio_util::sync::CancellationToken;
17
18use crate::image_processing::{
19    detect_image_mime, get_image_dimensions, process_image, ImageProcessingOptions, OutputFormat,
20};
21
22/// Input for the image tool.
23#[derive(Debug, Clone, JsonSchema, Deserialize)]
24pub struct ImageToolInput {
25    /// The action to perform: "info", "resize", "convert", or "process".
26    pub action: String,
27    /// Base64-encoded image data (required for all actions).
28    pub image_data: String,
29    /// Target width in pixels (for resize/process).
30    #[serde(default)]
31    pub width: Option<u32>,
32    /// Target height in pixels (for resize/process).
33    #[serde(default)]
34    pub height: Option<u32>,
35    /// Output format: "png", "jpeg", or "webp" (for convert/process).
36    #[serde(default)]
37    pub format: Option<String>,
38    /// JPEG/WebP quality 1-100 (for convert/process).
39    #[serde(default)]
40    pub quality: Option<u8>,
41    /// Whether to strip EXIF metadata (default: true).
42    #[serde(default)]
43    pub strip_exif: Option<bool>,
44}
45
46/// Image tool implementation.
47pub struct ImageTool {
48    schema: Tool,
49    max_dimension: u32,
50    default_format: String,
51}
52
53impl ImageTool {
54    /// Create a new image tool.
55    pub fn new(max_dimension: u32, default_format: &str) -> Self {
56        Self {
57            schema: Self::build_schema(),
58            max_dimension,
59            default_format: default_format.to_string(),
60        }
61    }
62
63    fn build_schema() -> Tool {
64        let params = schemars::schema_for!(ImageToolInput);
65        Tool {
66            name: "image".to_string(),
67            description: "Process, resize, convert, and get information about images. \
68                         Supports PNG, JPEG, WebP, GIF, and BMP formats. \
69                         Actions: 'info' (get metadata), 'resize' (scale dimensions), \
70                         'convert' (change format), 'process' (resize + convert + strip EXIF)."
71                .to_string(),
72            parameters: Schema::new(serde_json::to_value(params).unwrap_or_default()),
73            constrained_sampling: Some(ConstrainedSamplingConfig::JsonSchema {
74                strict: ConstrainedStrictness::Prefer,
75            }),
76        }
77    }
78
79    fn decode_base64(data: &str) -> Result<Vec<u8>, AgentError> {
80        STANDARD
81            .decode(data)
82            .map_err(|e| AgentError::tool(format!("Invalid base64: {e}")))
83    }
84
85    fn encode_base64(data: &[u8]) -> String {
86        STANDARD.encode(data)
87    }
88
89    fn parse_format(fmt: &str, quality: Option<u8>) -> Result<OutputFormat, AgentError> {
90        match fmt {
91            "png" => Ok(OutputFormat::Png),
92            "jpeg" | "jpg" => Ok(OutputFormat::Jpeg {
93                quality: quality.unwrap_or(85),
94            }),
95            "webp" => Ok(OutputFormat::WebP {
96                quality: quality.unwrap_or(85),
97            }),
98            other => Err(AgentError::tool(format!(
99                "Unsupported format: {other}. Use png, jpeg, or webp."
100            ))),
101        }
102    }
103
104    fn handle_info(&self, input: &ImageToolInput) -> Result<AgentToolResult, AgentError> {
105        let bytes = Self::decode_base64(&input.image_data)?;
106        let mime = detect_image_mime(&bytes).unwrap_or("unknown");
107        let (width, height) = get_image_dimensions(&bytes)
108            .map_err(|e| AgentError::tool(format!("Failed to read image: {e}")))?;
109
110        let output = serde_json::json!({
111            "success": true,
112            "width": width,
113            "height": height,
114            "mime_type": mime,
115            "size_bytes": bytes.len(),
116        });
117
118        Ok(AgentToolResult {
119            content: vec![TextContentOrImage::text(format!(
120                "Image info: {}x{} {} ({} bytes)",
121                width,
122                height,
123                mime,
124                bytes.len()
125            ))],
126            details: output,
127            usage: None,
128            added_tool_names: vec![],
129            terminate: false,
130        })
131    }
132
133    fn handle_resize(&self, input: &ImageToolInput) -> Result<AgentToolResult, AgentError> {
134        let bytes = Self::decode_base64(&input.image_data)?;
135
136        let options = ImageProcessingOptions {
137            max_width: input.width.map(|w| w.min(self.max_dimension)),
138            max_height: input.height.map(|h| h.min(self.max_dimension)),
139            format: None,
140            strip_exif: input.strip_exif.unwrap_or(true),
141        };
142
143        let processed = process_image(&bytes, options)
144            .map_err(|e| AgentError::tool(format!("Resize failed: {e}")))?;
145
146        let encoded = Self::encode_base64(&processed);
147        let (width, height) = get_image_dimensions(&processed).unwrap_or((0, 0));
148
149        let output = serde_json::json!({
150            "success": true,
151            "image_data": encoded,
152            "width": width,
153            "height": height,
154            "size_bytes": processed.len(),
155        });
156
157        Ok(AgentToolResult {
158            content: vec![TextContentOrImage::text(format!(
159                "Image resized to {}x{} ({} bytes)",
160                width,
161                height,
162                processed.len()
163            ))],
164            details: output,
165            usage: None,
166            added_tool_names: vec![],
167            terminate: false,
168        })
169    }
170
171    fn handle_convert(&self, input: &ImageToolInput) -> Result<AgentToolResult, AgentError> {
172        let bytes = Self::decode_base64(&input.image_data)?;
173        let format_str = input.format.as_deref().unwrap_or(&self.default_format);
174        let format = Self::parse_format(format_str, input.quality)?;
175
176        let options = ImageProcessingOptions {
177            max_width: None,
178            max_height: None,
179            format: Some(format),
180            strip_exif: input.strip_exif.unwrap_or(true),
181        };
182
183        let processed = process_image(&bytes, options)
184            .map_err(|e| AgentError::tool(format!("Convert failed: {e}")))?;
185
186        let encoded = Self::encode_base64(&processed);
187        let mime = detect_image_mime(&processed).unwrap_or("unknown");
188
189        let output = serde_json::json!({
190            "success": true,
191            "image_data": encoded,
192            "mime_type": mime,
193            "size_bytes": processed.len(),
194        });
195
196        Ok(AgentToolResult {
197            content: vec![TextContentOrImage::text(format!(
198                "Image converted to {} ({} bytes)",
199                mime,
200                processed.len()
201            ))],
202            details: output,
203            usage: None,
204            added_tool_names: vec![],
205            terminate: false,
206        })
207    }
208
209    fn handle_process(&self, input: &ImageToolInput) -> Result<AgentToolResult, AgentError> {
210        let bytes = Self::decode_base64(&input.image_data)?;
211
212        let format = if let Some(fmt) = &input.format {
213            Some(Self::parse_format(fmt, input.quality)?)
214        } else {
215            None
216        };
217
218        let options = ImageProcessingOptions {
219            max_width: input.width.map(|w| w.min(self.max_dimension)),
220            max_height: input.height.map(|h| h.min(self.max_dimension)),
221            format,
222            strip_exif: input.strip_exif.unwrap_or(true),
223        };
224
225        let processed = process_image(&bytes, options)
226            .map_err(|e| AgentError::tool(format!("Process failed: {e}")))?;
227
228        let encoded = Self::encode_base64(&processed);
229        let (width, height) = get_image_dimensions(&processed).unwrap_or((0, 0));
230        let mime = detect_image_mime(&processed).unwrap_or("unknown");
231
232        let output = serde_json::json!({
233            "success": true,
234            "image_data": encoded,
235            "width": width,
236            "height": height,
237            "mime_type": mime,
238            "size_bytes": processed.len(),
239        });
240
241        Ok(AgentToolResult {
242            content: vec![TextContentOrImage::text(format!(
243                "Image processed: {}x{} {} ({} bytes)",
244                width,
245                height,
246                mime,
247                processed.len()
248            ))],
249            details: output,
250            usage: None,
251            added_tool_names: vec![],
252            terminate: false,
253        })
254    }
255}
256
257#[async_trait]
258impl AgentTool for ImageTool {
259    fn schema(&self) -> &Tool {
260        &self.schema
261    }
262
263    fn label(&self) -> &str {
264        "image"
265    }
266
267    async fn execute(
268        &self,
269        _tool_call_id: &str,
270        params: serde_json::Value,
271        _signal: CancellationToken,
272        _on_update: Arc<dyn Fn(ToolResultPartial) + Send + Sync>,
273    ) -> Result<AgentToolResult, AgentError> {
274        let input: ImageToolInput = serde_json::from_value(params)
275            .map_err(|e| AgentError::tool(format!("image input invalid: {e}")))?;
276
277        match input.action.as_str() {
278            "info" => self.handle_info(&input),
279            "resize" => self.handle_resize(&input),
280            "convert" => self.handle_convert(&input),
281            "process" => self.handle_process(&input),
282            other => Err(AgentError::tool(format!(
283                "Unknown action: {other}. Use info, resize, convert, or process."
284            ))),
285        }
286    }
287}
288
289/// Create an image tool with default settings.
290pub fn create_image_tool() -> Arc<dyn AgentTool> {
291    Arc::new(ImageTool::new(2048, "png"))
292}
293
294/// Create an image tool with custom settings.
295pub fn create_image_tool_with_config(
296    max_dimension: u32,
297    default_format: &str,
298) -> Arc<dyn AgentTool> {
299    Arc::new(ImageTool::new(max_dimension, default_format))
300}
301
302/// Configuration type for the image tool (max_dimension, default_format).
303pub type ImageToolConfig = (u32, String);
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308
309    #[test]
310    fn test_parse_format() {
311        assert!(matches!(
312            ImageTool::parse_format("png", None),
313            Ok(OutputFormat::Png)
314        ));
315        assert!(matches!(
316            ImageTool::parse_format("jpeg", Some(90)),
317            Ok(OutputFormat::Jpeg { quality: 90 })
318        ));
319        assert!(matches!(
320            ImageTool::parse_format("webp", Some(80)),
321            Ok(OutputFormat::WebP { quality: 80 })
322        ));
323        assert!(ImageTool::parse_format("bmp", None).is_err());
324    }
325
326    #[test]
327    fn test_decode_base64() {
328        // Valid base64
329        assert!(ImageTool::decode_base64("aGVsbG8=").is_ok());
330        // Invalid base64
331        assert!(ImageTool::decode_base64("not-valid-base64!!!").is_err());
332    }
333
334    #[test]
335    fn test_encode_decode_roundtrip() {
336        let original = b"test image data";
337        let encoded = ImageTool::encode_base64(original);
338        let decoded = ImageTool::decode_base64(&encoded).unwrap();
339        assert_eq!(original.to_vec(), decoded);
340    }
341}
rpi_tools/tools/image.rs

rpi_tools/tools/
image.rs
