diff --git a/Cargo.lock b/Cargo.lock index 0845b98315c1..4421f13b7a83 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4679,6 +4679,7 @@ version = "1.37.0" dependencies = [ "anyhow", "async-trait", + "base64 0.22.1", "chrono", "fs-err", "futures", @@ -4693,6 +4694,7 @@ dependencies = [ "sec1", "serde", "serde_json", + "tempfile", "test-case", "thiserror 1.0.69", "tokio", diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index 348bfa476e67..e53c11b69058 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -40,3 +40,7 @@ fs-err = { version = "3.1", default-features = false } futures = { workspace = true } uuid = { workspace = true, features = ["v4"] } utoipa.workspace = true +base64.workspace = true + +[dev-dependencies] +tempfile.workspace = true diff --git a/crates/goose-providers/src/image.rs b/crates/goose-providers/src/image.rs index 8c2e2df18c75..9b04063b4e91 100644 --- a/crates/goose-providers/src/image.rs +++ b/crates/goose-providers/src/image.rs @@ -1,5 +1,9 @@ +use crate::errors::ProviderError; +use base64::Engine; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +use std::io::Read; +use std::path::Path; #[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq)] pub enum ImageFormat { @@ -26,6 +30,78 @@ pub fn convert_image(data: &str, mime_type: &str, image_format: &ImageFormat) -> } } +#[derive(Debug)] +pub struct EncodedImage { + pub mime_type: String, + pub data: String, +} + +fn is_image_file(path: &Path) -> bool { + if let Ok(mut file) = std::fs::File::open(path) { + let mut buffer = [0u8; 8]; + if file.read(&mut buffer).is_ok() { + return matches!( + &buffer[0..4], + [0x89, 0x50, 0x4E, 0x47] | [0xFF, 0xD8, 0xFF, _] | [0x47, 0x49, 0x46, 0x38] + ); + } + } + false +} + +pub fn detect_image_path(text: &str) -> Option<&str> { + let extensions = [".png", ".jpg", ".jpeg"]; + + for word in text.split_whitespace() { + if extensions + .iter() + .any(|ext| word.to_lowercase().ends_with(ext)) + { + let path = Path::new(word); + if path.is_absolute() && path.is_file() && is_image_file(path) { + return Some(word); + } + } + } + None +} + +pub fn load_image_file(path: &str) -> Result { + let path = Path::new(path); + + if !is_image_file(path) { + return Err(ProviderError::RequestFailed( + "File is not a valid image".to_string(), + )); + } + + let bytes = std::fs::read(path).map_err(|error| { + ProviderError::RequestFailed(format!("Failed to read image file: {error}")) + })?; + + let mime_type = match path.extension().and_then(|extension| extension.to_str()) { + Some(extension) => match extension.to_lowercase().as_str() { + "png" => "image/png", + "jpg" | "jpeg" => "image/jpeg", + _ => { + return Err(ProviderError::RequestFailed( + "Unsupported image format".to_string(), + )) + } + }, + None => { + return Err(ProviderError::RequestFailed( + "Unknown image format".to_string(), + )) + } + }; + + Ok(EncodedImage { + mime_type: mime_type.to_string(), + data: base64::prelude::BASE64_STANDARD.encode(&bytes), + }) +} + #[cfg(test)] mod tests { use super::*; @@ -38,6 +114,77 @@ mod tests { assert_eq!(decoded, ImageFormat::OpenAi); } + fn write_file(path: &std::path::Path, data: &[u8]) { + std::fs::write(path, data).unwrap(); + } + + #[test] + fn detects_absolute_image_paths() { + let temp_dir = tempfile::tempdir().unwrap(); + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, + ]; + write_file(&png_path, &png_data); + let png_path_str = png_path.to_str().unwrap(); + + let fake_png_path = temp_dir.path().join("fake.png"); + write_file(&fake_png_path, b"not a real png"); + + let text = format!("Here is an image {png_path_str}"); + assert_eq!(detect_image_path(&text), Some(png_path_str)); + + let text = format!("Here is a fake image {}", fake_png_path.to_str().unwrap()); + assert_eq!(detect_image_path(&text), None); + + assert_eq!( + detect_image_path("Here is a fake.png that doesn't exist"), + None + ); + assert_eq!(detect_image_path("Here is a file.txt"), None); + assert_eq!(detect_image_path("Here is a relative/path/image.png"), None); + } + + #[test] + fn loads_supported_image_files() { + let temp_dir = tempfile::tempdir().unwrap(); + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, + ]; + write_file(&png_path, &png_data); + + let image = load_image_file(png_path.to_str().unwrap()).unwrap(); + assert_eq!(image.mime_type, "image/png"); + assert!(!image.data.is_empty()); + } + + #[test] + fn rejects_invalid_or_unsupported_image_files() { + let temp_dir = tempfile::tempdir().unwrap(); + let fake_png_path = temp_dir.path().join("fake.png"); + write_file(&fake_png_path, b"not a real png"); + let result = load_image_file(fake_png_path.to_str().unwrap()); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("not a valid image")); + + let result = load_image_file("nonexistent.png"); + assert!(result.is_err()); + + let gif_path = temp_dir.path().join("test.gif"); + let gif_data = [0x47, 0x49, 0x46, 0x38, 0x39, 0x61]; + write_file(&gif_path, &gif_data); + let result = load_image_file(gif_path.to_str().unwrap()); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Unsupported image format")); + } + #[test] fn converts_image_to_openai_payload() { assert_eq!( diff --git a/crates/goose/src/providers/utils.rs b/crates/goose/src/providers/utils.rs index 568690639d9f..9cc5e4de969d 100644 --- a/crates/goose/src/providers/utils.rs +++ b/crates/goose/src/providers/utils.rs @@ -2,7 +2,6 @@ use super::base::Usage; use crate::config::paths::Paths; use crate::providers::errors::ProviderError; use anyhow::Result; -use base64::Engine; pub use goose_providers::function_name::{is_valid_function_name, sanitize_function_name}; pub use goose_providers::google_response::{handle_response_google_compat, is_google_model}; pub use goose_providers::image::ImageFormat; @@ -17,97 +16,20 @@ pub use goose_types::{ use rmcp::model::{AnnotateAble, ImageContent, RawImageContent}; use serde::Serialize; use serde_json::Value; -use std::io::Read; -use std::path::Path; /// Convert an image content into an image json based on format pub fn convert_image(image: &ImageContent, image_format: &ImageFormat) -> Value { goose_providers::image::convert_image(&image.data, &image.mime_type, image_format) } -/// Check if a file is actually an image by examining its magic bytes -fn is_image_file(path: &Path) -> bool { - if let Ok(mut file) = std::fs::File::open(path) { - let mut buffer = [0u8; 8]; // Large enough for most image magic numbers - if file.read(&mut buffer).is_ok() { - // Check magic numbers for common image formats - return match &buffer[0..4] { - // PNG: 89 50 4E 47 - [0x89, 0x50, 0x4E, 0x47] => true, - // JPEG: FF D8 FF - [0xFF, 0xD8, 0xFF, _] => true, - // GIF: 47 49 46 38 - [0x47, 0x49, 0x46, 0x38] => true, - _ => false, - }; - } - } - false -} - -/// Detect if a string contains a path to an image file -pub fn detect_image_path(text: &str) -> Option<&str> { - // Basic image file extension check - let extensions = [".png", ".jpg", ".jpeg"]; - - // Find any word that ends with an image extension - for word in text.split_whitespace() { - if extensions - .iter() - .any(|ext| word.to_lowercase().ends_with(ext)) - { - let path = Path::new(word); - // Check if it's an absolute path and file exists - if path.is_absolute() && path.is_file() { - // Verify it's actually an image file - if is_image_file(path) { - return Some(word); - } - } - } - } - None -} +pub use goose_providers::image::detect_image_path; /// Convert a local image file to base64 encoded ImageContent pub fn load_image_file(path: &str) -> Result { - let path = Path::new(path); - - // Verify it's an image before proceeding - if !is_image_file(path) { - return Err(ProviderError::RequestFailed( - "File is not a valid image".to_string(), - )); - } - - // Read the file - let bytes = std::fs::read(path) - .map_err(|e| ProviderError::RequestFailed(format!("Failed to read image file: {}", e)))?; - - // Detect mime type from extension - let mime_type = match path.extension().and_then(|e| e.to_str()) { - Some(ext) => match ext.to_lowercase().as_str() { - "png" => "image/png", - "jpg" | "jpeg" => "image/jpeg", - _ => { - return Err(ProviderError::RequestFailed( - "Unsupported image format".to_string(), - )) - } - }, - None => { - return Err(ProviderError::RequestFailed( - "Unknown image format".to_string(), - )) - } - }; - - // Convert to base64 - let data = base64::prelude::BASE64_STANDARD.encode(&bytes); - + let image = goose_providers::image::load_image_file(path)?; Ok(RawImageContent { - mime_type: mime_type.to_string(), - data, + mime_type: image.mime_type, + data: image.data, meta: None, } .no_annotation())