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(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,7 +11,7 @@ use llm_engine::tool::{Tool, ToolDefinition, ToolMeta};
|
||||
use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tools::{Tracker, core_builtin_tools};
|
||||
use tools::{Tracker, core_builtin_tools, view_image_tool};
|
||||
use workdir::{LocalWorkdirSession, WorkdirSessionHandle};
|
||||
|
||||
fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope {
|
||||
@@ -100,6 +100,31 @@ fn meta_has_description_and_schema() {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn view_image_reads_scoped_bytes_without_serializing_them_as_text() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let spill = TempDir::new().unwrap();
|
||||
let scope = scope_with_spill(dir.path(), spill.path());
|
||||
let session: WorkdirSessionHandle =
|
||||
Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf()));
|
||||
let png = b"\x89PNG\r\n\x1a\nprivate-image-body";
|
||||
std::fs::write(dir.path().join("image.png"), png).unwrap();
|
||||
let definition = view_image_tool(session);
|
||||
let (_meta, tool) = definition();
|
||||
|
||||
let output = call(&tool, json!({ "path": "image.png" })).await;
|
||||
assert_eq!(output.attachments.len(), 1);
|
||||
let llm_engine::tool::Attachment::Image(image) = &output.attachments[0];
|
||||
assert_eq!(image.mime_type(), "image/png");
|
||||
assert_eq!(image.data(), png);
|
||||
let serialized = serde_json::to_string(&output).unwrap();
|
||||
assert!(!serialized.contains("private-image-body"));
|
||||
assert!(!serialized.contains("attachments"));
|
||||
|
||||
let escaped = call_err(&tool, json!({ "path": "../outside.png" })).await;
|
||||
assert!(escaped.to_string().contains("scope") || escaped.to_string().contains("path"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_then_edit_then_read_roundtrip() {
|
||||
let (dir, _spill, reg) = setup();
|
||||
|
||||
Reference in New Issue
Block a user