feat: add scoped multimodal image attachments
This commit is contained in:
@@ -96,7 +96,11 @@ impl Tool for BashTool {
|
||||
} else {
|
||||
Some(output.content)
|
||||
};
|
||||
Ok(ToolOutput { summary, content })
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content,
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -85,6 +85,7 @@ impl Tool for EditTool {
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(preview),
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -69,6 +69,7 @@ impl Tool for GlobTool {
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: (!body.is_empty()).then_some(body),
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,6 +120,7 @@ impl Tool for GrepTool {
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: (!result.output.is_empty()).then_some(result.output),
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@ mod edit;
|
||||
mod glob;
|
||||
mod grep;
|
||||
mod read;
|
||||
mod view_image;
|
||||
mod web;
|
||||
mod write;
|
||||
|
||||
@@ -27,6 +28,7 @@ pub use glob::glob_tool;
|
||||
pub use grep::grep_tool;
|
||||
pub use read::read_tool;
|
||||
pub use tracker::Tracker;
|
||||
pub use view_image::view_image_tool;
|
||||
pub use web::{web_fetch_tool, web_search_tool};
|
||||
pub use write::write_tool;
|
||||
|
||||
|
||||
@@ -86,6 +86,7 @@ impl Tool for ReadTool {
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(rendered.body),
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
//! `ViewImage` tool — attach a bounded image from the scoped Workdir.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::tool::{
|
||||
Attachment, ImageAttachment, Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle};
|
||||
|
||||
use crate::error::ToolsError;
|
||||
|
||||
/// Maximum image body accepted for one model request.
|
||||
pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
|
||||
|
||||
const DESCRIPTION: &str = "Attach an image from the bound Workdir to the next model request. \
|
||||
The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP.";
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct ViewImageParams {
|
||||
/// Logical path relative to the bound Workdir root.
|
||||
path: String,
|
||||
}
|
||||
|
||||
struct ViewImageTool {
|
||||
session: WorkdirSessionHandle,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ViewImageTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let input: ViewImageParams = serde_json::from_str(input_json).map_err(|error| {
|
||||
ToolError::InvalidArgument(format!("invalid ViewImage input: {error}"))
|
||||
})?;
|
||||
let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?;
|
||||
let result = self
|
||||
.session
|
||||
.read(ReadRequest {
|
||||
path: path.clone(),
|
||||
offset: 0,
|
||||
limit: usize::MAX,
|
||||
// The scoped provider enforces this cap while reading, rather
|
||||
// than allocating an unbounded binary body first.
|
||||
max_bytes: MAX_IMAGE_BYTES + 1,
|
||||
})
|
||||
.await
|
||||
.map_err(ToolsError::from)?;
|
||||
|
||||
if result.truncated || result.bytes.len() > MAX_IMAGE_BYTES {
|
||||
return Err(ToolError::InvalidArgument(format!(
|
||||
"image exceeds the {MAX_IMAGE_BYTES}-byte limit"
|
||||
)));
|
||||
}
|
||||
let mime_type = detect_image_mime(&result.bytes).ok_or_else(|| {
|
||||
ToolError::InvalidArgument(
|
||||
"unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(),
|
||||
)
|
||||
})?;
|
||||
let bytes = result.bytes.len();
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary: format!("Attached image {path} ({mime_type}, {bytes} bytes)"),
|
||||
content: None,
|
||||
attachments: vec![Attachment::Image(ImageAttachment::new(
|
||||
mime_type,
|
||||
Arc::<[u8]>::from(result.bytes),
|
||||
))],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(ViewImageParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("ViewImage")
|
||||
.description(DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ViewImageTool {
|
||||
session: session.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn detect_image_mime(bytes: &[u8]) -> Option<&'static str> {
|
||||
if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
|
||||
Some("image/png")
|
||||
} else if bytes.starts_with(&[0xff, 0xd8, 0xff]) {
|
||||
Some("image/jpeg")
|
||||
} else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
|
||||
Some("image/gif")
|
||||
} else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
|
||||
Some("image/webp")
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detects_supported_image_signatures_without_trusting_extensions() {
|
||||
assert_eq!(
|
||||
detect_image_mime(b"\x89PNG\r\n\x1a\nbody"),
|
||||
Some("image/png")
|
||||
);
|
||||
assert_eq!(
|
||||
detect_image_mime(&[0xff, 0xd8, 0xff, 0xe0]),
|
||||
Some("image/jpeg")
|
||||
);
|
||||
assert_eq!(detect_image_mime(b"GIF89abody"), Some("image/gif"));
|
||||
assert_eq!(detect_image_mime(b"RIFF1234WEBPbody"), Some("image/webp"));
|
||||
assert_eq!(detect_image_mime(b"not an image"), None);
|
||||
}
|
||||
}
|
||||
@@ -1743,6 +1743,7 @@ fn json_output(value: Value) -> ToolOutput {
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(content),
|
||||
attachments: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -76,6 +76,7 @@ impl Tool for WriteTool {
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: None,
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user