124 lines
4.1 KiB
Rust
124 lines
4.1 KiB
Rust
//! `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);
|
|
}
|
|
}
|