feat: handle images mentioned in messages (#1202)

This commit is contained in:
Michael Neale
2025-02-12 14:36:46 -08:00
committed by GitHub
parent cfb1bba9b5
commit e461e6845d
2 changed files with 216 additions and 2 deletions
+51 -2
View File
@@ -3,7 +3,8 @@ use crate::model::ModelConfig;
use crate::providers::base::Usage;
use crate::providers::errors::ProviderError;
use crate::providers::utils::{
convert_image, is_valid_function_name, sanitize_function_name, ImageFormat,
convert_image, detect_image_path, is_valid_function_name, load_image_file,
sanitize_function_name, ImageFormat,
};
use anyhow::{anyhow, Error};
use mcp_core::ToolError;
@@ -26,7 +27,21 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
match content {
MessageContent::Text(text) => {
if !text.text.is_empty() {
converted["content"] = json!(text.text);
// Check for image paths in the text
if let Some(image_path) = detect_image_path(&text.text) {
// Try to load and convert the image
if let Ok(image) = load_image_file(image_path) {
converted["content"] = json!([
{"type": "text", "text": text.text},
convert_image(&image, image_format)
]);
} else {
// If image loading fails, just use the text
converted["content"] = json!(text.text);
}
} else {
converted["content"] = json!(text.text);
}
}
}
MessageContent::ToolRequest(request) => match &request.tool_call {
@@ -622,6 +637,40 @@ mod tests {
Ok(())
}
#[test]
fn test_format_messages_with_image_path() -> anyhow::Result<()> {
// Create a temporary PNG file with valid PNG magic numbers
let temp_dir = tempfile::tempdir()?;
let png_path = temp_dir.path().join("test.png");
let png_data = [
0x89, 0x50, 0x4E, 0x47, // PNG magic number
0x0D, 0x0A, 0x1A, 0x0A, // PNG header
0x00, 0x00, 0x00, 0x0D, // Rest of fake PNG data
];
std::fs::write(&png_path, &png_data)?;
let png_path_str = png_path.to_str().unwrap();
// Create message with image path
let message = Message::user().with_text(format!("Here is an image: {}", png_path_str));
let spec = format_messages(&[message], &ImageFormat::OpenAi);
assert_eq!(spec.len(), 1);
assert_eq!(spec[0]["role"], "user");
// Content should be an array with text and image
let content = spec[0]["content"].as_array().unwrap();
assert_eq!(content.len(), 2);
assert_eq!(content[0]["type"], "text");
assert!(content[0]["text"].as_str().unwrap().contains(png_path_str));
assert_eq!(content[1]["type"], "image_url");
assert!(content[1]["image_url"]["url"]
.as_str()
.unwrap()
.starts_with("data:image/png;base64,"));
Ok(())
}
#[test]
fn test_response_to_message_text() -> anyhow::Result<()> {
let response = json!({