microclaw_engine/tools/
generate_image.rs

1use async_trait::async_trait;
2use serde_json::{json, Value};
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use crate::internal::channels::channel::deliver_and_store_bot_message;
7use crate::internal::channels::channel_adapter::ChannelRegistry;
8use crate::internal::storage::db::Database;
9use crate::internal::tool_runtime::media_client::{persist_output, MediaClient};
10use crate::internal::tool_runtime::runtime::auth_context_from_input;
11use microclaw_core::llm_types::ToolDefinition;
12
13use super::{schema_object, Tool, ToolResult};
14use crate::config::{Config, ImageGenConfig, MediaConfig};
15
16const ALLOWED_SIZES: &[&str] = &[
17    "256x256",
18    "512x512",
19    "1024x1024",
20    "1024x1536",
21    "1536x1024",
22    "1024x1792",
23    "1792x1024",
24    "auto",
25];
26
27/// OpenAI-compatible image generation tool.
28///
29/// Calls `/v1/images/generations` on the configured base URL. Saves the
30/// returned PNG under `<data_dir>/media/images/<uuid>.png`. When the caller's
31/// channel supports `send_attachment`, the tool also delivers the image
32/// inline; otherwise the path is returned and the agent can hand it off
33/// manually.
34pub struct GenerateImageTool {
35    data_dir: PathBuf,
36    channels: Arc<ChannelRegistry>,
37    db: Arc<Database>,
38    cfg: ImageGenConfig,
39    media: MediaConfig,
40    openai_api_key: Option<String>,
41    openai_base_url: Option<String>,
42    timeout_secs: u64,
43}
44
45impl GenerateImageTool {
46    pub fn new(config: &Config, channels: Arc<ChannelRegistry>, db: Arc<Database>) -> Self {
47        Self {
48            data_dir: PathBuf::from(&config.data_dir),
49            channels,
50            db,
51            cfg: config.media.image_gen.clone(),
52            media: config.media.clone(),
53            openai_api_key: config.openai_api_key.clone(),
54            openai_base_url: config.openai_base_url.clone(),
55            timeout_secs: config.tool_timeout_secs("generate_image", 120),
56        }
57    }
58
59    fn client(&self) -> Result<MediaClient, String> {
60        let key = self
61            .media
62            .resolve_api_key(self.openai_api_key.as_deref())
63            .ok_or_else(|| {
64                "generate_image requires an API key. Set one of media.api_key, \
65                 MICROCLAW_OPENAI_API_KEY, OPENAI_API_KEY, or top-level openai_api_key."
66                    .to_string()
67            })?;
68        let base = self.media.resolve_base_url(self.openai_base_url.as_deref());
69        MediaClient::new(base, key, self.timeout_secs)
70    }
71}
72
73#[async_trait]
74impl Tool for GenerateImageTool {
75    fn name(&self) -> &str {
76        "generate_image"
77    }
78
79    fn definition(&self) -> ToolDefinition {
80        ToolDefinition {
81            name: self.name().into(),
82            description: "Generate an image from a natural-language prompt via an \
83                OpenAI-compatible images endpoint. The PNG is saved under the bot's \
84                data directory and — when the active channel supports attachments — \
85                also sent back to the user inline. Returns the file path and a \
86                confirmation of delivery. Disabled by default; requires operator \
87                opt-in via `media.image_gen.enabled`."
88                .into(),
89            input_schema: schema_object(
90                json!({
91                    "prompt": {
92                        "type": "string",
93                        "description": "What to draw. Be specific — style, subject, composition."
94                    },
95                    "size": {
96                        "type": "string",
97                        "enum": ALLOWED_SIZES,
98                        "description": "Optional image size. Defaults to media.image_gen.default_size."
99                    },
100                    "model": {
101                        "type": "string",
102                        "description": "Optional model override (e.g. gpt-image-1, dall-e-3)."
103                    },
104                    "deliver": {
105                        "type": "boolean",
106                        "description": "Whether to attempt channel delivery (default true). Set false to only return the path."
107                    }
108                }),
109                &["prompt"],
110            ),
111        }
112    }
113
114    async fn execute(&self, input: Value) -> ToolResult {
115        if !self.cfg.enabled {
116            return ToolResult::error(
117                "generate_image is disabled. Set media.image_gen.enabled=true to enable.".into(),
118            );
119        }
120        let prompt = match input.get("prompt").and_then(|v| v.as_str()) {
121            Some(p) if !p.trim().is_empty() => p.trim().to_string(),
122            _ => return ToolResult::error("Missing or empty parameter: prompt".into()),
123        };
124        let size = input
125            .get("size")
126            .and_then(|v| v.as_str())
127            .unwrap_or(self.cfg.default_size.as_str())
128            .to_string();
129        if !ALLOWED_SIZES.iter().any(|s| s == &size.as_str()) {
130            return ToolResult::error(format!(
131                "invalid size '{size}'. Allowed: {}",
132                ALLOWED_SIZES.join(", ")
133            ));
134        }
135        let model = input
136            .get("model")
137            .and_then(|v| v.as_str())
138            .map(|s| s.to_string())
139            .unwrap_or_else(|| self.cfg.model.clone());
140        let deliver = input
141            .get("deliver")
142            .and_then(|v| v.as_bool())
143            .unwrap_or(true);
144
145        let client = match self.client() {
146            Ok(c) => c,
147            Err(e) => return ToolResult::error(e),
148        };
149
150        let body = json!({
151            "prompt": prompt,
152            "size": size,
153            "model": model,
154            "n": 1,
155            // gpt-image-1 returns base64 by default; dall-e-3 returns a URL
156            // unless response_format is set. Ask for base64 explicitly so the
157            // return-shape is uniform.
158            "response_format": "b64_json",
159        });
160
161        let resp = match client
162            .post_json("images/generations")
163            .json(&body)
164            .send()
165            .await
166        {
167            Ok(r) => r,
168            Err(e) => return ToolResult::error(format!("images API request failed: {e}")),
169        };
170        if !resp.status().is_success() {
171            let status = resp.status();
172            let text = resp.text().await.unwrap_or_default();
173            return ToolResult::error(format!("images API HTTP {status}: {text}"));
174        }
175        let parsed: Value = match resp.json().await {
176            Ok(v) => v,
177            Err(e) => return ToolResult::error(format!("invalid JSON from images API: {e}")),
178        };
179
180        // Support both {"data":[{"b64_json":"..."}]} and
181        // {"data":[{"url":"..."}]}.
182        let datum = parsed
183            .get("data")
184            .and_then(|v| v.as_array())
185            .and_then(|arr| arr.first())
186            .cloned();
187        let bytes: Vec<u8> = if let Some(b64) = datum
188            .as_ref()
189            .and_then(|d| d.get("b64_json"))
190            .and_then(|v| v.as_str())
191        {
192            use base64::Engine;
193            match base64::engine::general_purpose::STANDARD.decode(b64) {
194                Ok(b) => b,
195                Err(e) => return ToolResult::error(format!("bad base64 in image response: {e}")),
196            }
197        } else if let Some(url) = datum
198            .as_ref()
199            .and_then(|d| d.get("url"))
200            .and_then(|v| v.as_str())
201        {
202            // Download the URL through the media client (so SSRF check applies).
203            let dl = match client
204                .client()
205                .get(url)
206                .send()
207                .await
208                .map_err(|e| e.to_string())
209            {
210                Ok(r) => r,
211                Err(e) => return ToolResult::error(format!("failed to download image URL: {e}")),
212            };
213            if !dl.status().is_success() {
214                return ToolResult::error(format!("image download HTTP {}", dl.status()));
215            }
216            match dl.bytes().await {
217                Ok(b) => b.to_vec(),
218                Err(e) => return ToolResult::error(format!("image body read failed: {e}")),
219            }
220        } else {
221            return ToolResult::error(
222                "images API response missing both b64_json and url fields".into(),
223            );
224        };
225
226        let saved = match persist_output(&self.data_dir, "images", "png", &bytes) {
227            Ok(p) => p,
228            Err(e) => return ToolResult::error(format!("failed to save image: {e}")),
229        };
230
231        let mut summary = format!(
232            "generated {} ({} bytes) -> {}",
233            model,
234            bytes.len(),
235            saved.display()
236        );
237
238        if deliver {
239            if let Some(auth) = auth_context_from_input(&input) {
240                match deliver_attachment(
241                    &self.channels,
242                    self.db.clone(),
243                    auth.caller_chat_id,
244                    &saved,
245                    &prompt,
246                )
247                .await
248                {
249                    Ok(msg) => summary.push_str(&format!("; {msg}")),
250                    Err(e) => summary.push_str(&format!("; delivery skipped: {e}")),
251                }
252            }
253        }
254
255        ToolResult::success(summary).with_metadata(json!({
256            "path": saved.to_string_lossy(),
257            "model": model,
258            "size": size,
259        }))
260    }
261}
262
263async fn deliver_attachment(
264    channels: &ChannelRegistry,
265    db: Arc<Database>,
266    chat_id: i64,
267    file: &std::path::Path,
268    caption: &str,
269) -> Result<String, String> {
270    let routing = match crate::internal::channels::channel::get_required_chat_routing(
271        channels,
272        db.clone(),
273        chat_id,
274    )
275    .await
276    {
277        Ok(r) => r,
278        Err(e) => return Err(e),
279    };
280    let Some(adapter) = channels.get(&routing.channel_name) else {
281        return Err(format!("no adapter for channel '{}'", routing.channel_name));
282    };
283    if adapter.is_local_only() {
284        return Ok(format!(
285            "channel '{}' is local-only, path retained at: {}",
286            routing.channel_name,
287            file.display()
288        ));
289    }
290    let external_chat_id = crate::internal::storage::db::call_blocking(db.clone(), move |d| {
291        d.get_chat_external_id(chat_id)
292    })
293    .await
294    .map_err(|e| e.to_string())?
295    .unwrap_or_else(|| chat_id.to_string());
296    let caption_short = caption.chars().take(120).collect::<String>();
297    match adapter
298        .send_attachment(&external_chat_id, file, Some(&caption_short))
299        .await
300    {
301        Ok(_) => {
302            // Store a chat-visible marker so transcripts reflect the send.
303            let _ = deliver_and_store_bot_message(
304                channels,
305                db,
306                "bot",
307                chat_id,
308                &format!("[image attached: {}]", file.display()),
309            )
310            .await;
311            Ok(format!("delivered via channel '{}'", routing.channel_name))
312        }
313        Err(e) => Err(e),
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    fn size_list_includes_defaults() {
323        assert!(ALLOWED_SIZES.contains(&"1024x1024"));
324        assert!(ALLOWED_SIZES.contains(&"auto"));
325    }
326}
microclaw_engine/tools/generate_image.rs

microclaw_engine/tools/
generate_image.rs
