From f4d290f5ed331b3296026808b49727a27277e338 Mon Sep 17 00:00:00 2001 From: Nick Pismenkov Date: Fri, 6 Mar 2026 17:55:41 -0800 Subject: [PATCH] feat: Support processing images by IronClaw --- .claude/settings.local.json | 13 ++ src/agent/dispatcher.rs | 26 ++++ src/agent/session.rs | 31 +++- src/agent/thread_ops.rs | 6 +- src/app.rs | 47 ++++++ src/channels/channel.rs | 12 ++ src/channels/repl.rs | 3 + src/channels/wasm/wrapper.rs | 5 + src/channels/web/mod.rs | 5 + src/channels/web/openai_compat.rs | 1 + src/channels/web/server.rs | 12 ++ src/channels/web/sse.rs | 1 + src/channels/web/static/app.js | 114 +++++++++++++- src/channels/web/static/index.html | 3 + src/channels/web/static/style.css | 98 ++++++++++++ src/channels/web/types.rs | 36 ++++- src/channels/web/ws.rs | 19 ++- src/llm/image_models.rs | 139 +++++++++++++++++ src/llm/mod.rs | 7 +- src/llm/nearai_chat.rs | 103 +++++++++---- src/llm/provider.rs | 29 ++++ src/llm/rig_adapter.rs | 34 ++++- src/llm/vision_models.rs | 160 ++++++++++++++++++++ src/tools/builtin/image_analyze.rs | 235 +++++++++++++++++++++++++++++ src/tools/builtin/image_edit.rs | 231 ++++++++++++++++++++++++++++ src/tools/builtin/image_gen.rs | 203 +++++++++++++++++++++++++ src/tools/builtin/mod.rs | 6 + src/tools/registry.rs | 41 ++++- tests/e2e_routine_heartbeat.rs | 3 + 29 files changed, 1572 insertions(+), 51 deletions(-) create mode 100644 .claude/settings.local.json create mode 100644 src/llm/image_models.rs create mode 100644 src/llm/vision_models.rs create mode 100644 src/tools/builtin/image_analyze.rs create mode 100644 src/tools/builtin/image_edit.rs create mode 100644 src/tools/builtin/image_gen.rs diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 00000000..38eda129 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,13 @@ +{ + "permissions": { + "allow": [ + "Bash(cargo check:*)", + "Bash(cargo clippy:*)", + "Bash(cargo test:*)", + "Bash(cargo fmt:*)", + "Bash(grep:*)", + "Bash(env:*)", + "Skill(ship)" + ] + } +} diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 95d8d711..6c2b518b 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -640,6 +640,32 @@ impl Agent { &message.metadata, ) .await; + + // Check for image_generated sentinel and emit SSE event + if let Ok(result_json) = + serde_json::from_str::(output) + && let Some("image_generated") = + result_json.get("type").and_then(|v| v.as_str()) + && let (Some(data), Some(media_type), Some(path)) = ( + result_json.get("data").and_then(|v| v.as_str()), + result_json.get("media_type").and_then(|v| v.as_str()), + result_json.get("path").and_then(|v| v.as_str()), + ) + { + let data_url = + format!("data:{};base64,{}", media_type, data); + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ImageGenerated { + data_url, + path: path.to_string(), + }, + &message.metadata, + ) + .await; + } } // Record result in thread diff --git a/src/agent/session.rs b/src/agent/session.rs index 4c3dbd67..db3cd9e4 100644 --- a/src/agent/session.rs +++ b/src/agent/session.rs @@ -16,7 +16,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::llm::{ChatMessage, ToolCall}; +use crate::llm::{ChatMessage, ImageAttachment, ToolCall}; /// A session containing one or more threads. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -250,6 +250,22 @@ impl Thread { &mut self.turns[turn_number] } + /// Start a new turn with user input and image attachments. + pub fn start_turn_with_images( + &mut self, + user_input: impl Into, + images: Vec, + ) -> &mut Turn { + let turn_number = self.turns.len(); + let mut turn = Turn::new(turn_number, user_input); + turn.images = images; + self.turns.push(turn); + self.state = ThreadState::Processing; + self.updated_at = Utc::now(); + // turn_number was len() before push, so it's a valid index after push + &mut self.turns[turn_number] + } + /// Complete the current turn with a response. pub fn complete_turn(&mut self, response: impl Into) { if let Some(turn) = self.turns.last_mut() { @@ -320,7 +336,14 @@ impl Thread { pub fn messages(&self) -> Vec { let mut messages = Vec::new(); for turn in &self.turns { - messages.push(ChatMessage::user(&turn.user_input)); + if turn.images.is_empty() { + messages.push(ChatMessage::user(&turn.user_input)); + } else { + messages.push(ChatMessage::user_with_images( + &turn.user_input, + turn.images.clone(), + )); + } if let Some(ref response) = turn.response { messages.push(ChatMessage::assistant(response)); } @@ -407,6 +430,9 @@ pub struct Turn { pub completed_at: Option>, /// Error message (if failed). pub error: Option, + /// Images attached to this turn's user input. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub images: Vec, } impl Turn { @@ -421,6 +447,7 @@ impl Turn { started_at: Utc::now(), completed_at: None, error: None, + images: Vec::new(), } } diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index bd1e5258..9ebad423 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -264,7 +264,11 @@ impl Agent { .threads .get_mut(&thread_id) .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; - thread.start_turn(content); + if message.images.is_empty() { + thread.start_turn(content); + } else { + thread.start_turn_with_images(content, message.images.clone()); + } thread.messages() }; diff --git a/src/app.rs b/src/app.rs index d273df41..efe2c7bd 100644 --- a/src/app.rs +++ b/src/app.rs @@ -376,6 +376,53 @@ impl AppBuilder { } let ws = Arc::new(ws); tools.register_memory_tools(Arc::clone(&ws)); + + // Register image tools if image generation models are available + match llm.list_models().await { + Ok(models) => { + if crate::llm::image_models::has_image_generation_model(&models) { + if let Some(image_model) = + crate::llm::image_models::suggest_image_model(&models) + { + tools.register_image_tools( + self.config.llm.nearai.clone(), + Arc::clone(&ws), + ); + tracing::info!( + "Image generation tools registered (model: {})", + image_model + ); + } + } else { + tracing::debug!( + "No image generation models detected in available models: {:?}", + models + ); + } + + // Register vision analysis tool if vision models are available + if crate::llm::vision_models::has_vision_model(&models) { + if let Some(vision_model) = + crate::llm::vision_models::suggest_vision_model(&models) + { + tools.register_vision_tools(Arc::clone(&ws)); + tracing::info!( + "Image analysis tool registered (vision model: {})", + vision_model + ); + } + } else { + tracing::debug!("No vision-capable models detected in available models"); + } + } + Err(e) => { + tracing::warn!( + "Failed to list available models for image tool registration: {}", + e + ); + } + } + Some(ws) } else { None diff --git a/src/channels/channel.rs b/src/channels/channel.rs index 46fbc9ca..2b3fd001 100644 --- a/src/channels/channel.rs +++ b/src/channels/channel.rs @@ -9,6 +9,7 @@ use futures::Stream; use uuid::Uuid; use crate::error::ChannelError; +use crate::llm::ImageAttachment; /// A message received from an external channel. #[derive(Debug, Clone)] @@ -29,6 +30,8 @@ pub struct IncomingMessage { pub received_at: DateTime, /// Channel-specific metadata. pub metadata: serde_json::Value, + /// Images attached to this message. + pub images: Vec, } impl IncomingMessage { @@ -47,6 +50,7 @@ impl IncomingMessage { thread_id: None, received_at: Utc::now(), metadata: serde_json::Value::Null, + images: Vec::new(), } } @@ -67,6 +71,12 @@ impl IncomingMessage { self.user_name = Some(name.into()); self } + + /// Attach image attachments. + pub fn with_images(mut self, images: Vec) -> Self { + self.images = images; + self + } } /// Stream of incoming messages. @@ -163,6 +173,8 @@ pub enum StatusUpdate { success: bool, message: String, }, + /// An image was generated or edited by a tool. + ImageGenerated { data_url: String, path: String }, } impl StatusUpdate { diff --git a/src/channels/repl.rs b/src/channels/repl.rs index f4cc7d3f..a7b5d108 100644 --- a/src/channels/repl.rs +++ b/src/channels/repl.rs @@ -585,6 +585,9 @@ impl Channel for ReplChannel { eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m"); } } + StatusUpdate::ImageGenerated { path, .. } => { + eprintln!(" \x1b[36m[image]\x1b[0m {path}"); + } } Ok(()) } diff --git a/src/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index 28272769..16304a08 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -2591,6 +2591,11 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha ), metadata_json, }, + StatusUpdate::ImageGenerated { path, .. } => wit_channel::StatusUpdate { + status: wit_channel::StatusType::Status, + message: format!("Image generated: {}", path), + metadata_json, + }, } } diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index 5152e551..aa292461 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -369,6 +369,11 @@ impl Channel for GatewayChannel { success, message, }, + StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated { + data_url, + path, + thread_id, + }, }; self.state.sse.broadcast(event); diff --git a/src/channels/web/openai_compat.rs b/src/channels/web/openai_compat.rs index c493ef5c..8e4c017f 100644 --- a/src/channels/web/openai_compat.rs +++ b/src/channels/web/openai_compat.rs @@ -247,6 +247,7 @@ pub fn convert_messages(messages: &[OpenAiMessage]) -> Result, tool_call_id: None, name: m.name.clone(), tool_calls: None, + images: Vec::new(), }), } }) diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index e456febd..44edb7cf 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -43,6 +43,7 @@ use crate::channels::web::types::*; use crate::channels::web::util::{build_turns_from_db_messages, truncate_preview}; use crate::db::Database; use crate::extensions::ExtensionManager; +use crate::llm::ImageAttachment; use crate::orchestrator::job_manager::ContainerJobManager; use crate::tools::ToolRegistry; use crate::workspace::Workspace; @@ -626,6 +627,17 @@ async fn chat_send_handler( msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id})); } + // Convert image data to ImageAttachment + let images: Vec = req + .images + .into_iter() + .map(|img| ImageAttachment { + media_type: img.media_type, + data: img.data, + }) + .collect(); + msg = msg.with_images(images); + let msg_id = msg.id; tracing::debug!( "[chat_send_handler] Created message id={}, content={:?}", diff --git a/src/channels/web/sse.rs b/src/channels/web/sse.rs index e1e2b270..809da048 100644 --- a/src/channels/web/sse.rs +++ b/src/channels/web/sse.rs @@ -143,6 +143,7 @@ impl SseManager { SseEvent::JobResult { .. } => "job_result", SseEvent::Heartbeat => "heartbeat", SseEvent::ExtensionStatus { .. } => "extension_status", + SseEvent::ImageGenerated { .. } => "image_generated", }; Ok(Event::default().event(event_type).data(data)) }); diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 0b69662d..1d1ff4ae 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -41,6 +41,9 @@ const SLASH_COMMANDS = [ let _slashSelected = -1; let _slashMatches = []; +// --- Image Attachments --- +let stagedImages = []; // Array of { media_type, data, previewUrl } + // --- Tool Activity State --- let _activeGroup = null; let _activeToolCards = {}; @@ -113,6 +116,78 @@ document.getElementById('token-input').addEventListener('keydown', (e) => { } })(); +// --- Image Attachment Handlers --- + +// Handle file picker selection +document.getElementById('image-input').addEventListener('change', (e) => { + const files = e.target.files; + if (files) { + for (let file of files) { + if (file.type.startsWith('image/')) { + const reader = new FileReader(); + reader.onload = (evt) => { + const base64Data = evt.target.result.split(',')[1]; // Remove data URL prefix + stagedImages.push({ + media_type: file.type, + data: base64Data, + previewUrl: evt.target.result, + }); + renderImagePreviews(); + }; + reader.readAsDataURL(file); + } + } + } + // Reset file input so the same file can be selected again + e.target.value = ''; +}); + +// Handle paste event +document.getElementById('chat-input').addEventListener('paste', (e) => { + const items = e.clipboardData.items; + for (let item of items) { + if (item.type.startsWith('image/')) { + e.preventDefault(); + const file = item.getAsFile(); + const reader = new FileReader(); + reader.onload = (evt) => { + const base64Data = evt.target.result.split(',')[1]; + stagedImages.push({ + media_type: item.type, + data: base64Data, + previewUrl: evt.target.result, + }); + renderImagePreviews(); + }; + reader.readAsDataURL(file); + } + } +}); + +function renderImagePreviews() { + const strip = document.getElementById('image-preview-strip'); + if (stagedImages.length === 0) { + strip.style.display = 'none'; + return; + } + strip.style.display = 'flex'; + strip.innerHTML = ''; + stagedImages.forEach((img, idx) => { + const container = document.createElement('div'); + container.className = 'image-preview'; + container.innerHTML = ` + Preview + + `; + strip.appendChild(container); + }); +} + +function removeImage(idx) { + stagedImages.splice(idx, 1); + renderImagePreviews(); +} + // --- API helper --- function apiFetch(path, options) { @@ -315,6 +390,12 @@ function connectSSE() { setToolCardOutput(data.name, data.preview); }); + eventSource.addEventListener('image_generated', (e) => { + const data = JSON.parse(e.data); + if (!isCurrentThread(data.thread_id)) return; + addGeneratedImage(data.data_url, data.path); + }); + eventSource.addEventListener('stream_chunk', (e) => { const data = JSON.parse(e.data); if (!isCurrentThread(data.thread_id)) return; @@ -430,19 +511,28 @@ function sendMessage() { return; } const content = input.value.trim(); - if (!content) return; + if (!content && stagedImages.length === 0) return; addMessage('user', content); input.value = ''; autoResizeTextarea(input); input.focus(); + const images = stagedImages.map(img => ({ + media_type: img.media_type, + data: img.data, + })); + apiFetch('/api/chat/send', { method: 'POST', - body: { content, thread_id: currentThreadId || undefined }, + body: { content, thread_id: currentThreadId || undefined, images }, }).catch((err) => { addMessage('system', 'Failed to send: ' + err.message); }); + + // Clear staged images after sending + stagedImages = []; + renderImagePreviews(); } function enableChatInput() { @@ -858,6 +948,26 @@ function finalizeActivityGroup() { _activeToolCards = {}; } +function addGeneratedImage(dataUrl, path) { + const container = document.getElementById('chat-messages'); + const card = document.createElement('div'); + card.className = 'generated-image-card'; + + const img = document.createElement('img'); + img.src = dataUrl; + img.alt = 'Generated image'; + img.className = 'generated-image'; + + const pathLabel = document.createElement('div'); + pathLabel.className = 'generated-image-path'; + pathLabel.textContent = 'Saved to: ' + path; + + card.appendChild(img); + card.appendChild(pathLabel); + container.appendChild(card); + container.scrollTop = container.scrollHeight; +} + function showApproval(data) { const container = document.getElementById('chat-messages'); const card = document.createElement('div'); diff --git a/src/channels/web/static/index.html b/src/channels/web/static/index.html index 1d232d17..382fdccd 100644 --- a/src/channels/web/static/index.html +++ b/src/channels/web/static/index.html @@ -129,7 +129,10 @@
+
+ +
diff --git a/src/channels/web/static/style.css b/src/channels/web/static/style.css index ead9cec8..7059c3d0 100644 --- a/src/channels/web/static/style.css +++ b/src/channels/web/static/style.css @@ -1093,6 +1093,37 @@ body { font-style: italic; } +/* Generated image card */ +.generated-image-card { + align-self: flex-start; + width: 50%; + background: var(--bg-secondary); + border: 1px solid var(--border); + border-radius: var(--radius-lg); + overflow: hidden; + margin: 8px 0; + box-shadow: var(--shadow); + display: flex; + flex-direction: column; + flex-shrink: 0; +} + +.generated-image { + display: block; + width: 100%; + border-radius: var(--radius-lg); + object-fit: contain; +} + +.generated-image-path { + padding: 8px 12px; + font-size: 12px; + color: var(--text-secondary); + background: var(--bg-tertiary); + border-top: 1px solid var(--border); + word-break: break-all; +} + /* Tool calls summary (persisted between user/assistant messages) */ .tool-calls-summary { background: var(--bg-secondary); @@ -1325,6 +1356,73 @@ body { cursor: not-allowed; } +.attach-btn { + padding: 8px 12px; + background: transparent; + color: var(--text-secondary); + border: 1px solid var(--border); + border-radius: var(--radius); + cursor: pointer; + font-size: 16px; + transition: all 0.2s; +} + +.attach-btn:hover { + background: var(--bg); + color: var(--text); + border-color: var(--accent); +} + +.image-preview-strip { + display: flex; + padding: 12px 16px 0 16px; + gap: 12px; + background: var(--bg-secondary); + overflow-x: auto; + border-top: 1px solid var(--border); +} + +.image-preview { + position: relative; + width: 80px; + height: 80px; + flex-shrink: 0; + border-radius: var(--radius); + overflow: hidden; + background: var(--bg); + border: 1px solid var(--border); +} + +.image-preview img { + width: 100%; + height: 100%; + object-fit: cover; +} + +.image-preview-remove { + position: absolute; + top: -1px; + right: -1px; + width: 24px; + height: 24px; + padding: 0; + background: rgba(0, 0, 0, 0.6); + color: white; + border: none; + border-radius: 0; + font-size: 18px; + font-weight: bold; + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + transition: background 0.2s; +} + +.image-preview-remove:hover { + background: rgba(0, 0, 0, 0.8); +} + /* Memory Tab */ .memory-container { flex: 1; diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 0e74e26e..145aa5e9 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -5,10 +5,18 @@ use uuid::Uuid; // --- Chat --- +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ImageData { + pub media_type: String, + pub data: String, // base64-encoded +} + #[derive(Debug, Deserialize)] pub struct SendMessageRequest { pub content: String, pub thread_id: Option, + #[serde(default)] + pub images: Vec, } #[derive(Debug, Serialize)] @@ -225,6 +233,17 @@ pub enum SseEvent { #[serde(skip_serializing_if = "Option::is_none")] message: Option, }, + + /// An image was generated or edited. + #[serde(rename = "image_generated")] + ImageGenerated { + /// Base64 data URL: "data:image/png;base64,..." + data_url: String, + /// Workspace path where the image is saved. + path: String, + #[serde(skip_serializing_if = "Option::is_none")] + thread_id: Option, + }, } // --- Memory --- @@ -606,6 +625,8 @@ pub enum WsClientMessage { Message { content: String, thread_id: Option, + #[serde(default)] + images: Vec, }, /// Approve or deny a pending tool execution. #[serde(rename = "approval")] @@ -673,6 +694,7 @@ impl WsServerMessage { SseEvent::JobStatus { .. } => "job_status", SseEvent::JobResult { .. } => "job_result", SseEvent::ExtensionStatus { .. } => "extension_status", + SseEvent::ImageGenerated { .. } => "image_generated", }; let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null); WsServerMessage::Event { @@ -791,9 +813,14 @@ mod tests { let json = r#"{"type":"message","content":"hello","thread_id":"t1"}"#; let msg: WsClientMessage = serde_json::from_str(json).unwrap(); match msg { - WsClientMessage::Message { content, thread_id } => { + WsClientMessage::Message { + content, + thread_id, + images, + } => { assert_eq!(content, "hello"); assert_eq!(thread_id.as_deref(), Some("t1")); + assert!(images.is_empty()); } _ => panic!("Expected Message variant"), } @@ -804,9 +831,14 @@ mod tests { let json = r#"{"type":"message","content":"hi"}"#; let msg: WsClientMessage = serde_json::from_str(json).unwrap(); match msg { - WsClientMessage::Message { content, thread_id } => { + WsClientMessage::Message { + content, + thread_id, + images, + } => { assert_eq!(content, "hi"); assert!(thread_id.is_none()); + assert!(images.is_empty()); } _ => panic!("Expected Message variant"), } diff --git a/src/channels/web/ws.rs b/src/channels/web/ws.rs index 2477217e..ac8e0e1c 100644 --- a/src/channels/web/ws.rs +++ b/src/channels/web/ws.rs @@ -22,6 +22,7 @@ use crate::agent::submission::Submission; use crate::channels::IncomingMessage; use crate::channels::web::server::GatewayState; use crate::channels::web::types::{WsClientMessage, WsServerMessage}; +use crate::llm::ImageAttachment; /// Tracks active WebSocket connections. pub struct WsConnectionTracker { @@ -156,12 +157,26 @@ async fn handle_client_message( direct_tx: &mpsc::Sender, ) { match msg { - WsClientMessage::Message { content, thread_id } => { + WsClientMessage::Message { + content, + thread_id, + images, + } => { let mut incoming = IncomingMessage::new("gateway", user_id, &content); if let Some(ref tid) = thread_id { incoming = incoming.with_thread(tid); } + // Convert image data to ImageAttachment + let image_attachments: Vec = images + .into_iter() + .map(|img| ImageAttachment { + media_type: img.media_type, + data: img.data, + }) + .collect(); + incoming = incoming.with_images(image_attachments); + let tx_guard = state.msg_tx.read().await; if let Some(ref tx) = *tx_guard { if tx.send(incoming).await.is_err() { @@ -349,6 +364,7 @@ mod tests { WsClientMessage::Message { content: "hello agent".to_string(), thread_id: Some("t1".to_string()), + images: vec![], }, &state, "user1", @@ -373,6 +389,7 @@ mod tests { WsClientMessage::Message { content: "hello".to_string(), thread_id: None, + images: vec![], }, &state, "user1", diff --git a/src/llm/image_models.rs b/src/llm/image_models.rs new file mode 100644 index 00000000..7714c51d --- /dev/null +++ b/src/llm/image_models.rs @@ -0,0 +1,139 @@ +//! Detection of image generation models across inference providers. + +/// Check if a model name indicates image generation capability. +/// +/// Detects models like: +/// - FLUX (Black Forest Labs): `flux`, `flux.2`, `flux-pro`, etc. +/// - DALL-E (OpenAI): `dall-e-2`, `dall-e-3`, etc. +/// - Stable Diffusion: `stable-diffusion`, `sdxl`, etc. +/// - Imagen (Google): `imagen`, `imagen-2`, etc. +/// - Other generation models +pub fn is_image_generation_model(model: &str) -> bool { + let model_lower = model.to_lowercase(); + + // FLUX models + if model_lower.contains("flux") { + return true; + } + + // DALL-E models + if model_lower.contains("dall-e") || model_lower.contains("dalle") { + return true; + } + + // Stable Diffusion models + if model_lower.contains("stable-diffusion") + || model_lower.contains("sdxl") + || model_lower.contains("stability") + { + return true; + } + + // Imagen models + if model_lower.contains("imagen") { + return true; + } + + // Midjourney (if exposed via API) + if model_lower.contains("midjourney") { + return true; + } + + // Replicate FLUX via API + if model_lower.contains("black-forest-labs") || model_lower.contains("lucataco") { + return true; + } + + false +} + +/// Check if any model in a list is an image generation model. +pub fn has_image_generation_model(models: &[String]) -> bool { + models.iter().any(|m| is_image_generation_model(m)) +} + +/// Suggest the best image generation model from available models. +/// +/// Priority: FLUX > DALL-E > others +pub fn suggest_image_model(models: &[String]) -> Option { + // Prefer FLUX + if let Some(flux) = models.iter().find(|m| m.to_lowercase().contains("flux")) { + return Some(flux.clone()); + } + + // Then DALL-E + if let Some(dalle) = models + .iter() + .find(|m| m.to_lowercase().contains("dall-e") || m.to_lowercase().contains("dalle")) + { + return Some(dalle.clone()); + } + + // Then any other image model + models + .iter() + .find(|m| is_image_generation_model(m)) + .cloned() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_flux_detection() { + assert!(is_image_generation_model( + "black-forest-labs/FLUX.2-klein-4B" + )); + assert!(is_image_generation_model("flux")); + assert!(is_image_generation_model("flux-pro")); + } + + #[test] + fn test_dalle_detection() { + assert!(is_image_generation_model("dall-e-3")); + assert!(is_image_generation_model("dall-e-2")); + assert!(is_image_generation_model("dalle-3")); + } + + #[test] + fn test_stable_diffusion_detection() { + assert!(is_image_generation_model("stable-diffusion-3")); + assert!(is_image_generation_model("sdxl")); + } + + #[test] + fn test_imagen_detection() { + assert!(is_image_generation_model("imagen")); + assert!(is_image_generation_model("imagen-3")); + } + + #[test] + fn test_non_image_models() { + assert!(!is_image_generation_model("claude-3-5-sonnet")); + assert!(!is_image_generation_model("gpt-4")); + assert!(!is_image_generation_model("gemini-pro")); + } + + #[test] + fn test_suggest_image_model() { + let models = vec![ + "gpt-4".to_string(), + "black-forest-labs/FLUX.2-klein-4B".to_string(), + "dall-e-3".to_string(), + ]; + + // Should prefer FLUX + assert_eq!( + suggest_image_model(&models), + Some("black-forest-labs/FLUX.2-klein-4B".to_string()) + ); + } + + #[test] + fn test_suggest_dalle_when_no_flux() { + let models = vec!["gpt-4".to_string(), "dall-e-3".to_string()]; + + assert_eq!(suggest_image_model(&models), Some("dall-e-3".to_string())); + } +} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 8ce4872a..6c994e30 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -10,6 +10,7 @@ pub mod circuit_breaker; pub mod costs; pub mod failover; +pub mod image_models; mod nearai_chat; mod provider; mod reasoning; @@ -19,13 +20,15 @@ pub mod retry; mod rig_adapter; pub mod session; pub mod smart_routing; +pub mod vision_models; pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider}; pub use failover::{CooldownConfig, FailoverProvider}; pub use nearai_chat::{ModelInfo, NearAiChatProvider}; pub use provider::{ - ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata, - Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult, + ChatMessage, CompletionRequest, CompletionResponse, FinishReason, ImageAttachment, LlmProvider, + ModelMetadata, Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, + ToolResult, }; pub use reasoning::{ ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN, diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 626c4d5c..e3762944 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -639,7 +639,7 @@ struct ChatCompletionRequest { struct ChatCompletionMessage { role: String, #[serde(skip_serializing_if = "Option::is_none")] - content: Option, + content: Option, #[serde(skip_serializing_if = "Option::is_none")] tool_call_id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -807,10 +807,15 @@ fn flatten_tool_messages(messages: Vec) -> Vec = Vec::new(); - if let Some(ref text) = msg.content - && !text.is_empty() - { - parts.push(text.clone()); + if let Some(content) = &msg.content { + // Extract string from JSON value + let text = match content { + serde_json::Value::String(s) => s.as_str(), + _ => "", + }; + if !text.is_empty() { + parts.push(text.to_string()); + } } for tc in calls { parts.push(format!( @@ -820,7 +825,7 @@ fn flatten_tool_messages(messages: Vec) -> Vec) -> Vec s.as_str(), + _ => "", + }; ChatCompletionMessage { role: "user".to_string(), - content: Some(format!("[Tool `{}` returned: {}]", tool_name, result)), + content: Some(serde_json::json!(format!( + "[Tool `{}` returned: {}]", + tool_name, result + ))), tool_call_id: None, name: None, @@ -870,8 +881,23 @@ impl From for ChatCompletionMessage { let content = if role == "assistant" && tool_calls.is_some() && msg.content.is_empty() { None + } else if !msg.images.is_empty() && role == "user" { + // User message with images: create a content array with text and image parts + let mut parts = vec![serde_json::json!({ + "type": "text", + "text": msg.content + })]; + for img in msg.images { + parts.push(serde_json::json!({ + "type": "image_url", + "image_url": { + "url": format!("data:{};base64,{}", img.media_type, img.data) + } + })); + } + Some(serde_json::Value::Array(parts)) } else { - Some(msg.content) + Some(serde_json::json!(msg.content)) }; Self { @@ -1038,7 +1064,7 @@ mod tests { let msg = ChatMessage::user("Hello"); let chat_msg: ChatCompletionMessage = msg.into(); assert_eq!(chat_msg.role, "user"); - assert_eq!(chat_msg.content, Some("Hello".to_string())); + assert_eq!(chat_msg.content, Some(serde_json::json!("Hello"))); } #[test] @@ -1112,14 +1138,14 @@ mod tests { let messages = vec![ ChatCompletionMessage { role: "system".to_string(), - content: Some("You are helpful.".to_string()), + content: Some(serde_json::json!("You are helpful.")), tool_call_id: None, name: None, tool_calls: None, }, ChatCompletionMessage { role: "user".to_string(), - content: Some("Hello".to_string()), + content: Some(serde_json::json!("Hello")), tool_call_id: None, name: None, tool_calls: None, @@ -1136,7 +1162,7 @@ mod tests { let messages = vec![ ChatCompletionMessage { role: "user".to_string(), - content: Some("test".to_string()), + content: Some(serde_json::json!("test")), tool_call_id: None, name: None, tool_calls: None, @@ -1157,7 +1183,7 @@ mod tests { }, ChatCompletionMessage { role: "tool".to_string(), - content: Some("hi".to_string()), + content: Some(serde_json::json!("hi")), tool_call_id: Some("call_1".to_string()), name: Some("echo".to_string()), tool_calls: None, @@ -1170,24 +1196,28 @@ mod tests { // Assistant tool_calls → plain assistant text assert_eq!(result[1].role, "assistant"); assert!(result[1].tool_calls.is_none()); - assert!( - result[1] - .content - .as_ref() - .unwrap() - .contains("[Called tool `echo`") - ); + if let Some(content) = &result[1].content { + if let serde_json::Value::String(s) = content { + assert!(s.contains("[Called tool `echo`")); + } else { + panic!("Content should be a string"); + } + } else { + panic!("Content should be present"); + } // Tool result → user message assert_eq!(result[2].role, "user"); assert!(result[2].tool_call_id.is_none()); - assert!( - result[2] - .content - .as_ref() - .unwrap() - .contains("[Tool `echo` returned: hi]") - ); + if let Some(content) = &result[2].content { + if let serde_json::Value::String(s) = content { + assert!(s.contains("[Tool `echo` returned: hi]")); + } else { + panic!("Content should be a string"); + } + } else { + panic!("Content should be present"); + } } #[test] @@ -1195,7 +1225,7 @@ mod tests { let messages = vec![ ChatCompletionMessage { role: "assistant".to_string(), - content: Some("Let me check that.".to_string()), + content: Some(serde_json::json!("Let me check that.")), tool_call_id: None, name: None, tool_calls: Some(vec![ChatCompletionToolCall { @@ -1209,7 +1239,7 @@ mod tests { }, ChatCompletionMessage { role: "tool".to_string(), - content: Some("found it".to_string()), + content: Some(serde_json::json!("found it")), tool_call_id: Some("call_1".to_string()), name: Some("search".to_string()), tool_calls: None, @@ -1217,9 +1247,16 @@ mod tests { ]; let result = flatten_tool_messages(messages); - let text = result[0].content.as_ref().unwrap(); - assert!(text.starts_with("Let me check that.")); - assert!(text.contains("[Called tool `search`")); + if let Some(content) = result[0].content.as_ref() { + if let serde_json::Value::String(text) = content { + assert!(text.starts_with("Let me check that.")); + assert!(text.contains("[Called tool `search`")); + } else { + panic!("Content should be a string"); + } + } else { + panic!("Content should be present"); + } } #[test] diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 84227df0..d8dcc858 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -16,6 +16,15 @@ pub enum Role { Tool, } +/// An image attachment for user messages. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageAttachment { + /// MIME type (e.g., "image/jpeg", "image/png", "image/gif", "image/webp") + pub media_type: String, + /// Base64-encoded image data (without data URL prefix) + pub data: String, +} + /// A message in a conversation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { @@ -31,6 +40,9 @@ pub struct ChatMessage { /// to appear on the assistant message preceding tool result messages). #[serde(skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, + /// Images attached to user messages. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub images: Vec, } impl ChatMessage { @@ -42,6 +54,7 @@ impl ChatMessage { tool_call_id: None, name: None, tool_calls: None, + images: Vec::new(), } } @@ -53,6 +66,19 @@ impl ChatMessage { tool_call_id: None, name: None, tool_calls: None, + images: Vec::new(), + } + } + + /// Create a user message with image attachments. + pub fn user_with_images(content: impl Into, images: Vec) -> Self { + Self { + role: Role::User, + content: content.into(), + tool_call_id: None, + name: None, + tool_calls: None, + images, } } @@ -64,6 +90,7 @@ impl ChatMessage { tool_call_id: None, name: None, tool_calls: None, + images: Vec::new(), } } @@ -82,6 +109,7 @@ impl ChatMessage { } else { Some(tool_calls) }, + images: Vec::new(), } } @@ -97,6 +125,7 @@ impl ChatMessage { tool_call_id: Some(tool_call_id.into()), name: Some(name.into()), tool_calls: None, + images: Vec::new(), } } } diff --git a/src/llm/rig_adapter.rs b/src/llm/rig_adapter.rs index da01b42c..f9b75efe 100644 --- a/src/llm/rig_adapter.rs +++ b/src/llm/rig_adapter.rs @@ -10,8 +10,8 @@ use rig::completion::{ ToolDefinition as RigToolDefinition, Usage as RigUsage, }; use rig::message::{ - Message as RigMessage, ToolChoice as RigToolChoice, ToolFunction, ToolResult as RigToolResult, - ToolResultContent, UserContent, + DocumentSourceKind, Image, ImageMediaType, Message as RigMessage, ToolChoice as RigToolChoice, + ToolFunction, ToolResult as RigToolResult, ToolResultContent, UserContent, }; use rust_decimal::Decimal; use serde::Serialize; @@ -230,7 +230,33 @@ fn convert_messages(messages: &[ChatMessage]) -> (Option, Vec { - history.push(RigMessage::user(&msg.content)); + if msg.images.is_empty() { + history.push(RigMessage::user(&msg.content)); + } else { + // User message with images: create multi-part content + let mut parts: Vec = vec![UserContent::text(&msg.content)]; + for img in &msg.images { + let media_type = match img.media_type.as_str() { + "image/jpeg" => ImageMediaType::JPEG, + "image/png" => ImageMediaType::PNG, + "image/gif" => ImageMediaType::GIF, + "image/webp" => ImageMediaType::WEBP, + _ => ImageMediaType::JPEG, + }; + parts.push(UserContent::Image(Image { + data: DocumentSourceKind::Base64(img.data.clone()), + media_type: Some(media_type), + detail: None, + additional_params: Default::default(), + })); + } + if let Ok(many) = OneOrMany::many(parts) { + history.push(RigMessage::User { content: many }); + } else { + // Fallback to text only + history.push(RigMessage::user(&msg.content)); + } + } } crate::llm::Role::Assistant => { if let Some(ref tool_calls) = msg.tool_calls { @@ -635,6 +661,7 @@ mod tests { tool_call_id: None, name: Some("search".to_string()), tool_calls: None, + images: vec![], }]; let (_preamble, history) = convert_messages(&messages); match &history[0] { @@ -784,6 +811,7 @@ mod tests { tool_call_id: None, name: Some("search".to_string()), tool_calls: None, + images: vec![], }; let messages = vec![assistant_msg, tool_result_msg]; let (_preamble, history) = convert_messages(&messages); diff --git a/src/llm/vision_models.rs b/src/llm/vision_models.rs new file mode 100644 index 00000000..1eb2f9ed --- /dev/null +++ b/src/llm/vision_models.rs @@ -0,0 +1,160 @@ +//! Detection of vision-capable models across inference providers. + +/// Check if a model name indicates vision capability. +/// +/// Detects models like: +/// - Claude (Anthropic): `claude-opus`, `claude-sonnet`, etc. +/// - GPT (OpenAI): `gpt-4-vision`, `gpt-4-turbo`, `gpt-4o`, etc. +/// - Gemini (Google): `gemini-pro-vision`, `gemini-2.0-flash`, etc. +/// - Llama (Meta): `llama-2-vision`, etc. +/// - Other vision-capable models +pub fn is_vision_model(model: &str) -> bool { + let model_lower = model.to_lowercase(); + + // Claude models (Anthropic) + if model_lower.contains("claude") { + return true; + } + + // GPT-4 models with vision support + if (model_lower.contains("gpt-4") + || model_lower.contains("gpt-4o") + || model_lower.contains("gpt-4-turbo") + || model_lower.contains("gpt-4-vision")) + && !model_lower.contains("gpt-4-mini") + { + return true; + } + + // Gemini models + if model_lower.contains("gemini") { + return true; + } + + // Llava and other vision models + if model_lower.contains("llava") + || model_lower.contains("vision") + || model_lower.contains("multimodal") + { + return true; + } + + false +} + +/// Check if any model in a list is a vision-capable model. +pub fn has_vision_model(models: &[String]) -> bool { + models.iter().any(|m| is_vision_model(m)) +} + +/// Suggest the best vision model from available models. +/// +/// Priority: Claude > GPT-4 > Gemini > others +pub fn suggest_vision_model(models: &[String]) -> Option { + // Prefer Claude + if let Some(claude) = models.iter().find(|m| m.to_lowercase().contains("claude")) { + return Some(claude.clone()); + } + + // Then GPT-4 + if let Some(gpt4) = models + .iter() + .find(|m| m.to_lowercase().contains("gpt-4") && !m.to_lowercase().contains("gpt-4-mini")) + { + return Some(gpt4.clone()); + } + + // Then Gemini + if let Some(gemini) = models.iter().find(|m| m.to_lowercase().contains("gemini")) { + return Some(gemini.clone()); + } + + // Then any other vision model + models.iter().find(|m| is_vision_model(m)).cloned() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_claude_detection() { + assert!(is_vision_model("claude-opus-4-20250514")); + assert!(is_vision_model("claude-sonnet-4-20250514")); + assert!(is_vision_model("claude-haiku-3-5-sonnet")); + } + + #[test] + fn test_gpt4_detection() { + assert!(is_vision_model("gpt-4-turbo")); + assert!(is_vision_model("gpt-4o")); + assert!(is_vision_model("gpt-4-vision")); + assert!(is_vision_model("gpt-4-32k")); + } + + #[test] + fn test_gpt4_mini_not_vision() { + assert!(!is_vision_model("gpt-4-mini")); + } + + #[test] + fn test_gemini_detection() { + assert!(is_vision_model("gemini-pro-vision")); + assert!(is_vision_model("gemini-2.0-flash")); + assert!(is_vision_model("gemini-1.5-pro")); + } + + #[test] + fn test_llava_detection() { + assert!(is_vision_model("llava-1.6")); + assert!(is_vision_model("llava-v1-7b")); + } + + #[test] + fn test_multimodal_detection() { + assert!(is_vision_model("my-multimodal-model")); + assert!(is_vision_model("custom-vision-model")); + } + + #[test] + fn test_non_vision_models() { + assert!(!is_vision_model("text-davinci-3")); + assert!(!is_vision_model("llama-2-7b")); + assert!(!is_vision_model("mistral-7b")); + } + + #[test] + fn test_suggest_vision_model() { + let models = vec![ + "gpt-4-turbo".to_string(), + "claude-opus-4-20250514".to_string(), + "gemini-2.0-flash".to_string(), + ]; + + // Should prefer Claude + assert_eq!( + suggest_vision_model(&models), + Some("claude-opus-4-20250514".to_string()) + ); + } + + #[test] + fn test_suggest_gpt4_when_no_claude() { + let models = vec!["gpt-4-turbo".to_string(), "gemini-2.0-flash".to_string()]; + + assert_eq!( + suggest_vision_model(&models), + Some("gpt-4-turbo".to_string()) + ); + } + + #[test] + fn test_suggest_gemini_when_no_claude_or_gpt4() { + let models = vec!["gemini-2.0-flash".to_string(), "text-davinci-3".to_string()]; + + assert_eq!( + suggest_vision_model(&models), + Some("gemini-2.0-flash".to_string()) + ); + } +} diff --git a/src/tools/builtin/image_analyze.rs b/src/tools/builtin/image_analyze.rs new file mode 100644 index 00000000..8e18bfe3 --- /dev/null +++ b/src/tools/builtin/image_analyze.rs @@ -0,0 +1,235 @@ +//! Image analysis tool for vision-capable LLMs. +//! +//! Reads images from the workspace and prepares them for vision analysis. +//! The LLM can then analyze the image content based on the user's query. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use crate::context::JobContext; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; +use crate::workspace::Workspace; + +/// Tool for analyzing images using a vision-capable LLM. +pub struct ImageAnalyzeTool { + workspace: Arc, +} + +impl ImageAnalyzeTool { + /// Create a new image analysis tool. + pub fn new(workspace: Arc) -> Self { + Self { workspace } + } + + /// Infer media type from file extension. + fn infer_media_type(path: &str) -> &'static str { + if path.ends_with(".png") || path.ends_with(".b64") { + "image/png" + } else if path.ends_with(".jpg") || path.ends_with(".jpeg") { + "image/jpeg" + } else if path.ends_with(".gif") { + "image/gif" + } else if path.ends_with(".webp") { + "image/webp" + } else { + "image/png" // Default to PNG + } + } +} + +#[async_trait] +impl Tool for ImageAnalyzeTool { + fn name(&self) -> &str { + "image_analyze" + } + + fn description(&self) -> &str { + "Analyze an image using the LLM's vision capabilities. Provide the workspace path to the image and a question or prompt about what you want to know about the image." + } + + fn parameters_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Workspace path to the image (e.g., 'images/generated/abc123.b64')" + }, + "query": { + "type": "string", + "description": "What do you want to know about the image? (e.g., 'describe the objects in this image', 'is there text in this image?')" + } + }, + "required": ["path", "query"] + }) + } + + async fn execute(&self, params: Value, _ctx: &JobContext) -> Result { + let start = std::time::Instant::now(); + + // Parse parameters + let path = params + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing or invalid 'path' parameter".to_string()) + })? + .to_string(); + + let query = params + .get("query") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing or invalid 'query' parameter".to_string()) + })? + .to_string(); + + if query.is_empty() { + return Err(ToolError::InvalidParameters( + "Query cannot be empty".to_string(), + )); + } + + // Read image from workspace + let doc = self.workspace.read(&path).await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to read image from workspace: {}", e)) + })?; + + // Infer media type from path + let media_type = Self::infer_media_type(&path).to_string(); + + // Return the image data and query so the agent can include the image in its vision analysis + Ok(ToolOutput::success( + json!({ + "type": "image_analysis_ready", + "path": path, + "query": query, + "data": doc.content, + "media_type": media_type, + "instruction": format!("The user wants you to analyze this image with the following query: {}", query) + }), + start.elapsed(), + )) + } + + fn requires_approval(&self, _params: &Value) -> ApprovalRequirement { + // Image analysis is read-only, no approval needed + ApprovalRequirement::Never + } + + fn sensitive_params(&self) -> &[&str] { + &[] + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_infer_media_type_png() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.png"), + "image/png" + ); + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.b64"), + "image/png" + ); + } + + #[test] + fn test_infer_media_type_jpeg() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.jpg"), + "image/jpeg" + ); + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.jpeg"), + "image/jpeg" + ); + } + + #[test] + fn test_infer_media_type_gif() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.gif"), + "image/gif" + ); + } + + #[test] + fn test_infer_media_type_webp() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.webp"), + "image/webp" + ); + } + + #[test] + fn test_infer_media_type_default() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.unknown"), + "image/png" + ); + } + + #[test] + fn test_parameters_schema_required_fields() { + let schema = json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Workspace path to the image (e.g., 'images/generated/abc123.b64')" + }, + "query": { + "type": "string", + "description": "What do you want to know about the image? (e.g., 'describe the objects in this image', 'is there text in this image?')" + } + }, + "required": ["path", "query"] + }); + + assert_eq!(schema["type"], "object"); + assert!(schema["properties"]["path"].is_object()); + assert!(schema["properties"]["query"].is_object()); + assert_eq!(schema["required"], json!(["path", "query"])); + } + + #[test] + fn test_infer_media_type_uppercase_extension_defaults() { + // Uppercase extensions don't match the lowercase checks, so defaults to PNG + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.PNG"), + "image/png" // Defaults to PNG for unknown extension + ); + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/test.JPG"), + "image/png" // Defaults to PNG for unknown extension + ); + } + + #[test] + fn test_infer_media_type_nested_path() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/generated/2024-03-06/deep/nested/image.png"), + "image/png" + ); + } + + #[test] + fn test_infer_media_type_multiple_dots() { + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/my.test.image.png"), + "image/png" + ); + assert_eq!( + ImageAnalyzeTool::infer_media_type("images/file.backup.jpg"), + "image/jpeg" + ); + } +} diff --git a/src/tools/builtin/image_edit.rs b/src/tools/builtin/image_edit.rs new file mode 100644 index 00000000..eaad524c --- /dev/null +++ b/src/tools/builtin/image_edit.rs @@ -0,0 +1,231 @@ +//! Image editing tool using NEAR AI cloud-api (FLUX model). + +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use base64::Engine; +use secrecy::ExposeSecret; +use serde_json::{Value, json}; +use uuid::Uuid; + +use crate::config::NearAiConfig; +use crate::context::JobContext; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, ToolRateLimitConfig}; +use crate::workspace::Workspace; + +/// Tool for editing existing images using NEAR AI cloud-api (FLUX). +pub struct ImageEditTool { + config: NearAiConfig, + client: reqwest::Client, + workspace: Arc, +} + +impl ImageEditTool { + /// Create a new image editing tool. + pub fn new(config: NearAiConfig, workspace: Arc) -> Self { + Self { + config, + client: reqwest::Client::new(), + workspace, + } + } +} + +#[async_trait] +impl Tool for ImageEditTool { + fn name(&self) -> &str { + "image_edit" + } + + fn description(&self) -> &str { + "Edit an existing image using NEAR AI cloud-api (FLUX) by providing the workspace path and a description of changes. \ + Returns the edited image saved to the workspace." + } + + fn parameters_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Workspace path to the source image (e.g., 'images/generated/abc123.b64')" + }, + "prompt": { + "type": "string", + "description": "Description of the edits to apply (max 4000 characters)" + }, + "size": { + "type": "string", + "enum": ["1024x1024", "1792x1024", "1024x1792"], + "description": "Image dimensions. Default: 1024x1024" + } + }, + "required": ["path", "prompt"] + }) + } + + async fn execute(&self, params: Value, _ctx: &JobContext) -> Result { + let start = std::time::Instant::now(); + + // Parse parameters + let path = params + .get("path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing or invalid 'path' parameter".to_string()) + })? + .to_string(); + + let prompt = params + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing or invalid 'prompt' parameter".to_string()) + })? + .to_string(); + + if prompt.is_empty() { + return Err(ToolError::InvalidParameters( + "Prompt cannot be empty".to_string(), + )); + } + + if prompt.len() > 4000 { + return Err(ToolError::InvalidParameters(format!( + "Prompt exceeds 4000 character limit (got {})", + prompt.len() + ))); + } + + let size = params + .get("size") + .and_then(|v| v.as_str()) + .unwrap_or("1024x1024"); + + // Read base64 image data from workspace + let doc = self.workspace.read(&path).await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to read image from workspace: {}", e)) + })?; + + // Decode base64 to bytes + let image_bytes = base64::engine::general_purpose::STANDARD + .decode(&doc.content) + .map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to decode base64 image data: {}", e)) + })?; + + // Build multipart form + let form = reqwest::multipart::Form::new() + .text("model", "black-forest-labs/FLUX.2-klein-4B") + .part( + "image", + reqwest::multipart::Part::bytes(image_bytes).file_name("image.png"), + ) + .text("prompt", prompt.clone()) + .text("n", "1") + .text("size", size.to_string()) + .text("response_format", "b64_json"); + + // Call NEAR AI cloud-api edit endpoint + let endpoint = format!( + "{}/v1/images/edits", + self.config.base_url.trim_end_matches('/') + ); + + let auth_header = if let Some(api_key) = &self.config.api_key { + format!("Bearer {}", api_key.expose_secret()) + } else { + "Bearer ".to_string() + }; + + let response = self + .client + .post(&endpoint) + .header("Authorization", auth_header) + .multipart(form) + .timeout(Duration::from_secs(120)) + .send() + .await + .map_err(|e| ToolError::ExternalService(format!("NEAR AI image edit failed: {}", e)))?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response + .text() + .await + .unwrap_or_else(|_| "Unknown error".to_string()); + return Err(ToolError::ExternalService(format!( + "NEAR AI image edit error ({}): {}", + status, error_text + ))); + } + + let response_json: Value = response.json().await.map_err(|e| { + ToolError::ExternalService(format!("Failed to parse NEAR AI response: {}", e)) + })?; + + // Extract base64 edited image data + let edited_base64 = response_json + .get("data") + .and_then(|d| d.get(0)) + .and_then(|item| item.get("b64_json")) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExternalService( + "Invalid NEAR AI response structure: missing base64 data".to_string(), + ) + })? + .to_string(); + + // Generate unique filename for edited image + let edit_id = Uuid::new_v4().to_string(); + let filename = format!("images/generated/{}_edit.png", edit_id); + + // Store edited image to workspace + let edit_path = format!("images/generated/{}_edit.b64", edit_id); + self.workspace + .write(&edit_path, &edited_base64) + .await + .map_err(|e| { + ToolError::ExecutionFailed(format!( + "Failed to save edited image to workspace: {}", + e + )) + })?; + + // Return sentinel JSON for agent_loop to detect and emit SSE event + Ok(ToolOutput::success( + json!({ + "type": "image_generated", + "path": edit_path, + "data": edited_base64, + "media_type": "image/png", + "prompt": prompt, + "size": size, + "filename": filename, + "source_path": path + }), + start.elapsed(), + )) + } + + fn requires_approval(&self, _params: &Value) -> ApprovalRequirement { + // Image editing is read-only on external state + ApprovalRequirement::Never + } + + fn rate_limit_config(&self) -> Option { + // DALL-E is expensive; rate limit aggressively + Some(ToolRateLimitConfig::new(6, 30)) + } + + fn sensitive_params(&self) -> &[&str] { + &[] + } + + fn execution_timeout(&self) -> std::time::Duration { + // Image editing can take 2+ minutes on the NEAR AI cloud-api + std::time::Duration::from_secs(180) + } +} diff --git a/src/tools/builtin/image_gen.rs b/src/tools/builtin/image_gen.rs new file mode 100644 index 00000000..a3721243 --- /dev/null +++ b/src/tools/builtin/image_gen.rs @@ -0,0 +1,203 @@ +//! Image generation tool using NEAR AI cloud-api (FLUX model). + +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use secrecy::ExposeSecret; +use serde_json::{Value, json}; +use uuid::Uuid; + +use crate::config::NearAiConfig; +use crate::context::JobContext; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, ToolRateLimitConfig}; +use crate::workspace::Workspace; + +/// Tool for generating images from text prompts using NEAR AI cloud-api (FLUX). +pub struct ImageGenerateTool { + config: NearAiConfig, + client: reqwest::Client, + workspace: Arc, +} + +impl ImageGenerateTool { + /// Create a new image generation tool. + pub fn new(config: NearAiConfig, workspace: Arc) -> Self { + Self { + config, + client: reqwest::Client::new(), + workspace, + } + } +} + +#[async_trait] +impl Tool for ImageGenerateTool { + fn name(&self) -> &str { + "image_generate" + } + + fn description(&self) -> &str { + "Generate an image from a text prompt using NEAR AI cloud-api (FLUX.2-klein-4B). \ + Returns the generated image saved to the workspace." + } + + fn parameters_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Detailed text description of the image to generate (max 4000 characters)" + }, + "size": { + "type": "string", + "enum": ["1024x1024", "1792x1024", "1024x1792"], + "description": "Image dimensions. Default: 1024x1024" + } + }, + "required": ["prompt"] + }) + } + + async fn execute(&self, params: Value, _ctx: &JobContext) -> Result { + let start = std::time::Instant::now(); + + // Parse parameters + let prompt = params + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing or invalid 'prompt' parameter".to_string()) + })? + .to_string(); + + if prompt.is_empty() { + return Err(ToolError::InvalidParameters( + "Prompt cannot be empty".to_string(), + )); + } + + if prompt.len() > 4000 { + return Err(ToolError::InvalidParameters(format!( + "Prompt exceeds 4000 character limit (got {})", + prompt.len() + ))); + } + + let size = params + .get("size") + .and_then(|v| v.as_str()) + .unwrap_or("1024x1024"); + + // Call NEAR AI cloud-api for image generation (FLUX model) + let request_body = json!({ + "model": "black-forest-labs/FLUX.2-klein-4B", + "prompt": prompt, + "n": 1, + "size": size, + "response_format": "b64_json" + }); + + let endpoint = format!( + "{}/v1/images/generations", + self.config.base_url.trim_end_matches('/') + ); + + let auth_header = if let Some(api_key) = &self.config.api_key { + format!("Bearer {}", api_key.expose_secret()) + } else { + // Fallback: use default NEAR AI cloud-api without explicit key + // (expects auth via environment or other mechanism) + "Bearer ".to_string() + }; + + let response = self + .client + .post(&endpoint) + .header("Authorization", auth_header) + .json(&request_body) + .timeout(Duration::from_secs(120)) + .send() + .await + .map_err(|e| { + ToolError::ExternalService(format!("NEAR AI image generation failed: {}", e)) + })?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response + .text() + .await + .unwrap_or_else(|_| "Unknown error".to_string()); + return Err(ToolError::ExternalService(format!( + "NEAR AI image generation error ({}): {}", + status, error_text + ))); + } + + let response_json: Value = response.json().await.map_err(|e| { + ToolError::ExternalService(format!("Failed to parse NEAR AI response: {}", e)) + })?; + + // Extract base64 image data + let base64_data = response_json + .get("data") + .and_then(|d| d.get(0)) + .and_then(|item| item.get("b64_json")) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExternalService( + "Invalid NEAR AI response structure: missing base64 data".to_string(), + ) + })? + .to_string(); + + // Generate unique filename + let image_id = Uuid::new_v4().to_string(); + let filename = format!("images/generated/{}.png", image_id); + + // Store the image file (with extension) containing the base64 data + let image_path = format!("images/generated/{}.b64", image_id); + self.workspace + .write(&image_path, &base64_data) + .await + .map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to save image to workspace: {}", e)) + })?; + + // Return sentinel JSON for agent_loop to detect and emit SSE event + Ok(ToolOutput::success( + json!({ + "type": "image_generated", + "path": image_path, + "data": base64_data, + "media_type": "image/png", + "prompt": prompt, + "size": size, + "filename": filename + }), + start.elapsed(), + )) + } + + fn requires_approval(&self, _params: &Value) -> ApprovalRequirement { + // Image generation from a prompt is read-only on external state + // so no approval needed + ApprovalRequirement::Never + } + + fn rate_limit_config(&self) -> Option { + // DALL-E is expensive; rate limit aggressively + Some(ToolRateLimitConfig::new(6, 30)) + } + + fn sensitive_params(&self) -> &[&str] { + &[] + } + + fn execution_timeout(&self) -> std::time::Duration { + // Image generation can take 2+ minutes on the NEAR AI cloud-api + std::time::Duration::from_secs(180) + } +} diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index 4931e5b8..065fe378 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -4,6 +4,9 @@ mod echo; pub mod extension_tools; mod file; mod http; +mod image_analyze; +mod image_edit; +mod image_gen; mod job; mod json; mod memory; @@ -23,6 +26,9 @@ pub use extension_tools::{ }; pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool}; pub use http::HttpTool; +pub use image_analyze::ImageAnalyzeTool; +pub use image_edit::ImageEditTool; +pub use image_gen::ImageGenerateTool; pub use job::{ CancelJobTool, CreateJobTool, JobEventsTool, JobPromptTool, JobStatusTool, ListJobsTool, PromptQueue, SchedulerSlot, diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 5809305e..1a90e6b8 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -17,11 +17,11 @@ use crate::skills::registry::SkillRegistry; use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder}; use crate::tools::builtin::{ ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool, - JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool, - MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool, - ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool, - ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, - WriteFileTool, + ImageEditTool, ImageGenerateTool, JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, + ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, + PromptQueue, ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, + SkillSearchTool, TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, + ToolRemoveTool, ToolSearchTool, WriteFileTool, }; use crate::tools::rate_limiter::RateLimiter; use crate::tools::tool::{Tool, ToolDomain}; @@ -71,6 +71,9 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "message", "web_fetch", "restart", + "image_generate", + "image_edit", + "image_analyze", ]; /// Registry of available tools. @@ -302,6 +305,34 @@ impl ToolRegistry { tracing::info!("Registered 4 memory tools"); } + /// Register image generation tools with NEAR AI config and workspace. + /// + /// Image tools require NEAR AI cloud-api access and workspace for storing generated images. + pub fn register_image_tools( + &self, + config: crate::config::NearAiConfig, + workspace: Arc, + ) { + self.register_sync(Arc::new(ImageGenerateTool::new( + config.clone(), + Arc::clone(&workspace), + ))); + self.register_sync(Arc::new(ImageEditTool::new(config, workspace))); + + tracing::info!("Registered 2 image tools (NEAR AI FLUX)"); + } + + /// Register image analysis tool with workspace access. + /// + /// Vision tool allows analyzing images using the LLM's vision capabilities. + pub fn register_vision_tools(&self, workspace: Arc) { + self.register_sync(Arc::new(crate::tools::builtin::ImageAnalyzeTool::new( + workspace, + ))); + + tracing::info!("Registered 1 vision tool (image analysis)"); + } + /// Register job management tools. /// /// Job tools allow the LLM to create, list, check status, and cancel jobs. diff --git a/tests/e2e_routine_heartbeat.rs b/tests/e2e_routine_heartbeat.rs index 6f4dda34..74ab0c9d 100644 --- a/tests/e2e_routine_heartbeat.rs +++ b/tests/e2e_routine_heartbeat.rs @@ -203,6 +203,7 @@ mod tests { thread_id: None, received_at: Utc::now(), metadata: serde_json::json!({}), + images: vec![], }; let fired = engine.check_event_triggers(&matching_msg).await; assert!( @@ -223,6 +224,7 @@ mod tests { thread_id: None, received_at: Utc::now(), metadata: serde_json::json!({}), + images: vec![], }; let fired_neg = engine.check_event_triggers(&non_matching_msg).await; assert_eq!(fired_neg, 0, "Expected 0 routines fired on non-match"); @@ -286,6 +288,7 @@ mod tests { thread_id: None, received_at: Utc::now(), metadata: serde_json::json!({}), + images: vec![], }; let fired1 = engine.check_event_triggers(&msg).await; assert!(fired1 >= 1, "First fire should work");