From f6c46764c392dfd5f2a1d070c146fb1c3d0b5aad Mon Sep 17 00:00:00 2001 From: andber1 <82754113+andber1@users.noreply.github.com> Date: Sat, 1 Aug 2026 19:08:55 +0200 Subject: [PATCH] Add support for pdf and image file references --- Cargo.lock | 2 + Cargo.toml | 2 + README.md | 5 ++ src/chat.rs | 13 ++-- src/input.rs | 176 +++++++++++++++++++++++++++++++++++++++++-------- src/main.rs | 8 ++- src/session.rs | 95 ++++++++++++++++++++++---- 7 files changed, 255 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5da4e04..77a448d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2769,6 +2769,8 @@ dependencies = [ "grep", "ignore", "libc", + "mime", + "mime_guess", "openrouter-rs", "pulldown-cmark", "regex", diff --git a/Cargo.toml b/Cargo.toml index 503452c..656455b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,8 @@ fuzzy-matcher = "0.3" schemars = "1.1" similar = "2.4" glob = "0.3" +mime_guess = "2.0" +mime = "0.3" rustyline = "14.0" grep = "0.3" ignore = "0.4" diff --git a/README.md b/README.md index 3a644ec..625f334 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,7 @@ A fast, minimal CLI for interacting with LLMs via OpenRouter. - **Fuzzy model selection** - `/sonnet`, `/gpt4`, `/haiku` - **File references** - `@file.txt` to include, `!file.txt` to write, `@!file.txt` to read+write +- **Multimodal inputs** - `@file.pdf`, `@file.png`, `@file.jpg`, `@file.gif`, and `@file.webp` are sent as native attachments - **Interactive chat** - Multi-turn conversations with `--chat` - **STDIN piping** - `cat code.rs | zo "review this"` - **Streaming output** - Syntax-highlighted markdown in real-time @@ -23,6 +24,10 @@ zo /sonnet "Explain async/await" # Include files zo "@main.rs Review this code" +# Include a PDF or image +zo "@document.pdf Summarize this" +zo "@diagram.png Describe this image" + #$ Pipeline git diff | zo 'Summarize these changes' diff --git a/src/chat.rs b/src/chat.rs index d955fde..6701c5b 100644 --- a/src/chat.rs +++ b/src/chat.rs @@ -7,13 +7,14 @@ use anyhow::{Context, Result}; use openrouter_rs::OpenRouterClient; +use openrouter_rs::api::chat::Content; use std::io::{self, BufRead, Write}; use crate::config::InlineColors; use crate::input::parse_file_patterns; use crate::models::ModelEntry; use crate::readline::ChatReadline; -use crate::session::{Session, build_user_message}; +use crate::session::{Session, build_typed_user_message}; use crate::shell::ShellRuntime; use crate::tools::ToolAccess; @@ -72,7 +73,7 @@ pub async fn run_chat_session( .context("Failed to parse file patterns from initial prompt")?; // Build first message combining file references, prompt, and STDIN - let first_message = build_user_message( + let first_message = build_typed_user_message( &initial_file_refs, &final_initial_prompt, options.initial_stdin.as_deref(), @@ -105,7 +106,11 @@ pub async fn run_chat_session( println!("Type 'exit', 'quit', or press Ctrl+D to end the conversation.\n"); // Send first message only if there's content - let has_initial_content = !first_message.trim().is_empty(); + let has_initial_content = match &first_message.content { + Content::Text(content) => !content.trim().is_empty(), + Content::Parts(parts) => !parts.is_empty(), + _ => true, + }; if has_initial_content { match session.send_message(first_message).await { Ok(_) => { @@ -155,7 +160,7 @@ pub async fn run_chat_session( } // Build message with file references (use expanded prompt) - let message = build_user_message(&file_refs, &final_input_prompt, None); + let message = build_typed_user_message(&file_refs, &final_input_prompt, None); // Send message and get response match session.send_message(message).await { diff --git a/src/input.rs b/src/input.rs index cec0b11..34ebb9b 100644 --- a/src/input.rs +++ b/src/input.rs @@ -1,5 +1,7 @@ use anyhow::{Context, Result, bail}; +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use glob::glob; +use mime::Mime; use std::collections::HashMap; use std::fs; use std::io::{self, IsTerminal, Read}; @@ -246,14 +248,21 @@ fn resolve_file_pattern_impl(pattern: &str, allow_missing: bool) -> Result, } +fn attachment_mime(path: &str) -> Option<&'static str> { + match mime_guess::from_path(path) + .first() + .as_ref() + .map(Mime::essence_str) + { + Some("application/pdf") => Some("application/pdf"), + Some("image/png") => Some("image/png"), + Some("image/jpeg") => Some("image/jpeg"), + Some("image/gif") => Some("image/gif"), + Some("image/webp") => Some("image/webp"), + _ => None, + } +} + +/// Read an input file as text or an encoded PDF/image attachment. +/// +/// `is_input_output` is true for `@!` references, which may name a new text file. +fn read_input_file(filename: &str, is_input_output: bool) -> Result { + if let Some(attachment_mime) = attachment_mime(filename) { + if is_input_output { + bail!("Input/output file '{}' must be a text file", filename); + } + let bytes = fs::read(filename).with_context(|| { + format!( + "Could not read file '{}'. Make sure it exists and is readable.", + filename + ) + })?; + let data_url = format!("data:{attachment_mime};base64,{}", BASE64.encode(bytes)); + return Ok(FileReference { + filename: filename.to_string(), + kind: FileReferenceKind::Attachment { + mime_type: attachment_mime.to_string(), + data_url, + }, + }); + } + + let content = if is_input_output { + read_file_if_exists(filename)? + } else { + fs::read_to_string(filename).with_context(|| { + format!( + "Could not read file '{}'. Make sure the file exists and is readable as UTF-8.", + filename + ) + })? + }; + Ok(FileReference { + filename: filename.to_string(), + kind: FileReferenceKind::Text(content), + }) +} + /// Parse all file patterns (@, !, @!) in a single pass /// /// This unified parser handles: @@ -438,23 +502,10 @@ pub fn parse_file_patterns( // Build input file references (@ and @!) if pattern.syntax_type.is_input() { for filename in &pattern.resolved_files { - let content = if pattern.syntax_type == FileSyntaxType::InputOutput { - // @! syntax - file might not exist yet - read_file_if_exists(filename)? - } else { - // @ syntax - file must exist - fs::read_to_string(filename).with_context(|| { - format!( - "Could not read file '{}'. Make sure the file exists and is readable.", - filename - ) - })? - }; - - file_references.push(FileReference { - filename: filename.clone(), - content, - }); + file_references.push(read_input_file( + filename, + pattern.syntax_type == FileSyntaxType::InputOutput, + )?); } } @@ -778,12 +829,65 @@ mod tests { let (_prompt, refs, _outputs) = result.unwrap(); assert_eq!(refs.len(), 1); assert_eq!(refs[0].filename, "test_file_single.txt"); - assert_eq!(refs[0].content.trim(), "test content"); + assert!( + matches!(&refs[0].kind, FileReferenceKind::Text(content) if content.trim() == "test content") + ); // Cleanup std::fs::remove_file(temp_file).ok(); } + #[test] + fn test_parse_file_patterns_classifies_pdf() { + // Create a temporary file for testing + use std::io::Write; + let pdf = "test_attachment.pdf"; + let mut file = std::fs::File::create(pdf).unwrap(); + writeln!(file, "pdf").unwrap(); + + let result = parse_file_patterns("@test_attachment.pdf"); + assert!(result.is_ok()); + let (_, refs, _) = result.unwrap(); + assert_eq!(refs.len(), 1); + assert!(matches!( + &refs[0].kind, + FileReferenceKind::Attachment { mime_type, .. } if mime_type == "application/pdf" + )); + assert!( + matches!(&refs[0].kind, FileReferenceKind::Attachment { data_url, .. } + if data_url.starts_with("data:application/pdf;base64,")) + ); + + // Cleanup + std::fs::remove_file(pdf).ok(); + } + + #[test] + fn test_parse_file_patterns_classifies_image() { + // Create a temporary file for testing + use std::io::Write; + let image = "test_attachment.png"; + let mut file = std::fs::File::create(image).unwrap(); + writeln!(file, "png").unwrap(); + + let result = parse_file_patterns("@test_attachment.png"); + assert!(result.is_ok()); + let (_, refs, _) = result.unwrap(); + assert_eq!(refs.len(), 1); + assert!(matches!( + &refs[0].kind, + FileReferenceKind::Attachment { mime_type, .. } if mime_type == "image/png" + )); + assert!(matches!( + &refs[0].kind, + FileReferenceKind::Attachment { data_url, .. } + if data_url.starts_with("data:image/png;base64,") + )); + + // Cleanup + std::fs::remove_file(image).ok(); + } + #[test] fn test_parse_file_patterns_input_multiple_files() { // Create temporary files for testing @@ -813,8 +917,12 @@ mod tests { .find(|r| r.filename == "test_file_multi2.txt") .unwrap(); - assert_eq!(file1_ref.content.trim(), "content 1"); - assert_eq!(file2_ref.content.trim(), "content 2"); + assert!( + matches!(&file1_ref.kind, FileReferenceKind::Text(content) if content.trim() == "content 1") + ); + assert!( + matches!(&file2_ref.kind, FileReferenceKind::Text(content) if content.trim() == "content 2") + ); // Cleanup std::fs::remove_file(temp_file1).ok(); @@ -854,7 +962,9 @@ mod tests { let (_prompt, refs, _outputs) = result.unwrap(); assert_eq!(refs.len(), 1); assert_eq!(refs[0].filename, "test_file_end.txt"); - assert_eq!(refs[0].content.trim(), "ending content"); + assert!( + matches!(&refs[0].kind, FileReferenceKind::Text(content) if content.trim() == "ending content") + ); // Cleanup std::fs::remove_file(temp_file).ok(); @@ -983,7 +1093,9 @@ mod tests { // Should have one file reference (because @! includes input) assert_eq!(refs.len(), 1); assert_eq!(refs[0].filename, "test_output_ref.txt"); - assert_eq!(refs[0].content.trim(), "test content"); + assert!( + matches!(&refs[0].kind, FileReferenceKind::Text(content) if content.trim() == "test content") + ); // Should also have one output file assert_eq!(outputs.len(), 1); @@ -1027,8 +1139,12 @@ mod tests { .find(|r| r.filename == "test_output_ref2.txt") .unwrap(); - assert_eq!(input_ref.content.trim(), "input content"); - assert_eq!(output_ref.content.trim(), "output content"); + assert!( + matches!(&input_ref.kind, FileReferenceKind::Text(content) if content.trim() == "input content") + ); + assert!( + matches!(&output_ref.kind, FileReferenceKind::Text(content) if content.trim() == "output content") + ); // Should have one output file (only @! file) assert_eq!(outputs.len(), 1); @@ -1674,7 +1790,7 @@ mod tests { } #[cfg(test)] mod integration_test { - use crate::input::parse_input; + use crate::input::{FileReferenceKind, parse_input}; use std::fs; use std::io::Write; @@ -1712,7 +1828,9 @@ mod integration_test { "Should have 1 file reference" ); assert_eq!(parsed.file_references[0].filename, test_file); - assert_eq!(parsed.file_references[0].content.trim(), "existing content"); + assert!( + matches!(&parsed.file_references[0].kind, FileReferenceKind::Text(content) if content.trim() == "existing content") + ); // Cleanup fs::remove_file(test_file).ok(); diff --git a/src/main.rs b/src/main.rs index 464455b..dd2ae0e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -237,7 +237,11 @@ fn display_debug_info( if !parsed_input.file_references.is_empty() { println!("\nFile References:"); for file_ref in &parsed_input.file_references { - println!(" - {}", file_ref.filename); + if let input::FileReferenceKind::Attachment { mime_type, .. } = &file_ref.kind { + println!(" - {} ({})", file_ref.filename, mime_type); + } else { + println!(" - {} (text)", file_ref.filename); + } } } else { println!("\nFile References: (none)"); @@ -630,7 +634,7 @@ async fn main() -> Result<()> { } // Build user message - let user_message = session::build_user_message( + let user_message = session::build_typed_user_message( &parsed_input.file_references, &parsed_input.prompt, parsed_input.stdin_content.as_deref(), diff --git a/src/session.rs b/src/session.rs index 57a22b5..223c296 100644 --- a/src/session.rs +++ b/src/session.rs @@ -10,7 +10,7 @@ use anyhow::{Context, Result}; use futures_util::StreamExt; use openrouter_rs::OpenRouterClient; -use openrouter_rs::api::chat::{ChatCompletionRequest, Message}; +use openrouter_rs::api::chat::{ChatCompletionRequest, ContentPart, Message}; use openrouter_rs::types::completion::FinishReason; use openrouter_rs::types::stream::StreamEvent; use openrouter_rs::types::typed_tool::TypedTool; @@ -23,7 +23,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::config::InlineColors; use crate::file_ops::FileWriter; -use crate::input::{FileReference, OutputFileSpec}; +use crate::input::{FileReference, FileReferenceKind, OutputFileSpec}; use crate::models::ModelEntry; use crate::render::StreamRenderer; use crate::shell::{RunProgramParams, RunShellCommandParams, ShellRuntime}; @@ -151,7 +151,7 @@ impl Session { } /// Send a user message and get the assistant's response. - pub async fn send_message(&mut self, user_content: String) -> Result { + pub async fn send_message(&mut self, user_message: Message) -> Result { if self.messages.is_empty() { let system_prompt = self.build_system_prompt(); if !system_prompt.is_empty() { @@ -160,8 +160,7 @@ impl Session { } } - self.messages - .push(Message::new(Role::User, user_content.as_str())); + self.messages.push(user_message); let mut all_response_text = String::new(); @@ -682,11 +681,12 @@ pub fn build_user_message( if !file_references.is_empty() { let files_formatted = file_references .iter() - .map(|file_ref| { - format!( + .filter_map(|file_ref| match &file_ref.kind { + FileReferenceKind::Text(content) => Some(format!( "\n{}\n", - file_ref.filename, file_ref.content - ) + file_ref.filename, content + )), + FileReferenceKind::Attachment { .. } => None, }) .collect::>() .join("\n\n"); @@ -704,6 +704,53 @@ pub fn build_user_message( parts.join("\n\n") } +/// Build the typed user message used for text and multimodal requests. +pub fn build_typed_user_message( + file_references: &[FileReference], + prompt: &str, + stdin_content: Option<&str>, +) -> Message { + let attachments = file_references + .iter() + .filter_map(|file_ref| match &file_ref.kind { + FileReferenceKind::Attachment { + mime_type, + data_url, + } => Some(( + file_ref.filename.as_str(), + mime_type.as_str(), + data_url.as_str(), + )), + FileReferenceKind::Text(_) => None, + }); + + if !file_references + .iter() + .any(|file_ref| matches!(&file_ref.kind, FileReferenceKind::Attachment { .. })) + { + return Message::new( + Role::User, + build_user_message(file_references, prompt, stdin_content), + ); + } + + let mut parts = Vec::new(); + let text = build_user_message(file_references, prompt, stdin_content); + if !text.is_empty() { + parts.push(ContentPart::text(text)); + } + + for (filename, mime_type, data_url) in attachments { + if mime_type == "application/pdf" { + parts.push(ContentPart::file_data_with_filename(data_url, filename)); + } else { + parts.push(ContentPart::image_url(data_url)); + } + } + + Message::with_parts(Role::User, parts) +} + #[cfg(test)] mod tests { use super::*; @@ -749,7 +796,7 @@ mod tests { fn test_build_user_message_with_file() { let files = vec![FileReference { filename: "test.txt".to_string(), - content: "content".to_string(), + kind: FileReferenceKind::Text("content".to_string()), }]; let msg = build_user_message(&files, "check this", None); assert!(msg.contains("")); @@ -761,7 +808,7 @@ mod tests { fn test_build_user_message_all_parts() { let files = vec![FileReference { filename: "data.csv".to_string(), - content: "col1,col2".to_string(), + kind: FileReferenceKind::Text("col1,col2".to_string()), }]; let msg = build_user_message(&files, "analyze", Some("extra data")); assert!(msg.contains("")); @@ -776,6 +823,32 @@ mod tests { assert_eq!(msg, "piped input"); } + #[test] + fn test_text_message_serializes_content_as_string() { + let message = build_typed_user_message(&[], "hello", None); + let value = serde_json::to_value(message).unwrap(); + assert!(value["content"].is_string()); + } + + #[test] + fn test_multipart_message_serializes_attachment_parts() { + let attachment = FileReference { + filename: "doc.pdf".to_string(), + kind: FileReferenceKind::Attachment { + mime_type: "application/pdf".to_string(), + data_url: "data:application/pdf;base64,ZmFrZQ==".to_string(), + }, + }; + let message = build_typed_user_message(&[attachment], "summarize", None); + let value = serde_json::to_value(message).unwrap(); + assert!(value["content"].is_array()); + assert_eq!(value["content"][1]["file"]["filename"], "doc.pdf"); + assert_eq!( + value["content"][1]["file"]["file_data"], + "data:application/pdf;base64,ZmFrZQ==" + ); + } + #[test] fn test_tool_availability_disabled_no_outputs() { let availability = determine_tool_availability(