use lopdf::{content::Content as PdfContent, Document, Object}; use mcp_core::{Content, ToolError}; use std::{fs, path::Path}; pub async fn pdf_tool( path: &str, operation: &str, cache_dir: &Path, ) -> Result, ToolError> { // Open and parse the PDF file let doc = Document::load(path) .map_err(|e| ToolError::ExecutionError(format!("Failed to open PDF file: {}", e)))?; let result = match operation { "extract_text" => { let mut text = String::new(); // Iterate over each page in the document for (page_num, page_id) in doc.get_pages() { text.push_str(&format!("Page {}:\n", page_num)); // Try to get text from page contents if let Ok(page_obj) = doc.get_object(page_id) { if let Ok(page_dict) = page_obj.as_dict() { // Try to get text from Contents stream if let Ok(contents) = page_dict.get(b"Contents").and_then(|c| c.as_reference()) { if let Ok(content_obj) = doc.get_object(contents) { if let Ok(stream) = content_obj.as_stream() { if let Ok(content_data) = stream.get_plain_content() { if let Ok(content) = PdfContent::decode(&content_data) { // Process each operation in the content stream for operation in content.operations { match operation.operator.as_ref() { // "Tj" operator: show text "Tj" => { for operand in operation.operands { if let Object::String(ref bytes, _) = operand { if let Ok(s) = std::str::from_utf8(bytes) { text.push_str(s); } } } text.push(' '); } // "TJ" operator: show text with positioning "TJ" => { if let Some(Object::Array(ref arr)) = operation.operands.first() { let mut last_was_text = false; for element in arr { match element { Object::String( ref bytes, _, ) => { if let Ok(s) = std::str::from_utf8( bytes, ) { if last_was_text { text.push(' '); } text.push_str(s); last_was_text = true; } } Object::Integer(offset) => { // Large negative offsets often indicate word spacing if *offset < -100 { text.push(' '); last_was_text = false; } } Object::Real(offset) => { if *offset < -100.0 { text.push(' '); last_was_text = false; } } _ => {} } } text.push(' '); } } _ => (), // Ignore other operators } } } } } } } } } text.push('\n'); } if text.trim().is_empty() { "No text found in PDF".to_string() } else { format!("Extracted text from PDF:\n\n{}", text) } } "extract_images" => { let cache_dir = cache_dir.join("pdf_images"); fs::create_dir_all(&cache_dir).map_err(|e| { ToolError::ExecutionError(format!("Failed to create image cache directory: {}", e)) })?; let mut images = Vec::new(); let mut image_count = 0; // Helper function to determine file extension based on stream dict fn get_image_extension(dict: &lopdf::Dictionary) -> &'static str { if let Ok(filter) = dict.get(b"Filter") { match filter { Object::Name(name) => { match name.as_slice() { b"DCTDecode" => ".jpg", b"JBIG2Decode" => ".jbig2", b"JPXDecode" => ".jp2", b"CCITTFaxDecode" => ".tiff", b"FlateDecode" => { // PNG-like images often use FlateDecode // Check color space to confirm if let Ok(cs) = dict.get(b"ColorSpace") { if let Ok(name) = cs.as_name() { if name == b"DeviceRGB" || name == b"DeviceGray" { return ".png"; } } } ".raw" } _ => ".raw", } } Object::Array(filters) => { // If multiple filters, check the last one if let Some(Object::Name(name)) = filters.last() { match name.as_slice() { b"DCTDecode" => return ".jpg", b"JPXDecode" => return ".jp2", _ => {} } } ".raw" } _ => ".raw", } } else { ".raw" } } // Process each page for (page_num, page_id) in doc.get_pages() { let page = doc.get_object(page_id).map_err(|e| { ToolError::ExecutionError(format!("Failed to get page {}: {}", page_num, e)) })?; let page_dict = page.as_dict().map_err(|e| { ToolError::ExecutionError(format!( "Failed to get page dict {}: {}", page_num, e )) })?; // Get page resources - handle both direct dict and reference let resources = match page_dict.get(b"Resources") { Ok(res) => match res { Object::Dictionary(dict) => Ok(dict), Object::Reference(id) => doc .get_object(*id) .map_err(|e| { ToolError::ExecutionError(format!( "Failed to get resource reference: {}", e )) }) .and_then(|obj| { obj.as_dict().map_err(|e| { ToolError::ExecutionError(format!( "Resource reference is not a dictionary: {}", e )) }) }), _ => Err(ToolError::ExecutionError( "Resources is neither dictionary nor reference".to_string(), )), }, Err(e) => Err(ToolError::ExecutionError(format!( "Failed to get Resources: {}", e ))), }?; // Look for XObject dictionary - handle both direct dict and reference let xobjects = match resources.get(b"XObject") { Ok(xobj) => match xobj { Object::Dictionary(dict) => Ok(dict), Object::Reference(id) => doc .get_object(*id) .map_err(|e| { ToolError::ExecutionError(format!( "Failed to get XObject reference: {}", e )) }) .and_then(|obj| { obj.as_dict().map_err(|e| { ToolError::ExecutionError(format!( "XObject reference is not a dictionary: {}", e )) }) }), _ => Err(ToolError::ExecutionError( "XObject is neither dictionary nor reference".to_string(), )), }, Err(e) => Err(ToolError::ExecutionError(format!( "Failed to get XObject: {}", e ))), }; if let Ok(xobjects) = xobjects { for (name, xobject) in xobjects.iter() { let xobject_id = xobject.as_reference().map_err(|_| { ToolError::ExecutionError("Failed to get XObject reference".to_string()) })?; let xobject = doc.get_object(xobject_id).map_err(|e| { ToolError::ExecutionError(format!("Failed to get XObject: {}", e)) })?; if let Ok(stream) = xobject.as_stream() { // Check if it's an image if let Ok(subtype) = stream.dict.get(b"Subtype").and_then(|s| s.as_name()) { if subtype == b"Image" { let extension = get_image_extension(&stream.dict); // Get image metadata let width = stream .dict .get(b"Width") .and_then(|w| w.as_i64()) .unwrap_or(0); let height = stream .dict .get(b"Height") .and_then(|h| h.as_i64()) .unwrap_or(0); let bpc = stream .dict .get(b"BitsPerComponent") .and_then(|b| b.as_i64()) .unwrap_or(0); // Get the image data if let Ok(data) = stream.get_plain_content() { let image_path = cache_dir.join(format!( "page{}_obj{}_{}{}", page_num, xobject_id.0, String::from_utf8_lossy(name), extension )); fs::write(&image_path, &data).map_err(|e| { ToolError::ExecutionError(format!( "Failed to write image: {}", e )) })?; images.push(format!( "Saved image to: {} ({}x{}, {} bits per component)", image_path.display(), width, height, bpc )); image_count += 1; } } } } } } } if images.is_empty() { "No images found in PDF".to_string() } else { format!("Found {} images:\n{}", image_count, images.join("\n")) } } _ => { return Err(ToolError::InvalidParameters(format!( "Invalid operation: {}. Valid operations are: 'extract_text', 'extract_images'", operation ))) } }; Ok(vec![Content::text(result)]) } #[cfg(test)] mod tests { use super::*; use std::path::PathBuf; #[tokio::test] async fn test_pdf_text_extraction() { let test_pdf_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("src/computercontroller/tests/data/test.pdf"); let cache_dir = tempfile::tempdir().unwrap().into_path(); println!("Testing text extraction from: {}", test_pdf_path.display()); let result = pdf_tool(test_pdf_path.to_str().unwrap(), "extract_text", &cache_dir).await; assert!(result.is_ok(), "PDF text extraction should succeed"); let content = result.unwrap(); assert!(!content.is_empty(), "Extracted text should not be empty"); let text = content[0].as_text().unwrap(); println!("Extracted text:\n{}", text); assert!(text.contains("Page 1"), "Should contain page marker"); assert!( text.contains("This is a test PDF"), "Should contain expected test content" ); } #[tokio::test] async fn test_pdf_image_extraction() { let test_pdf_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("src/computercontroller/tests/data/test_image.pdf"); let cache_dir = tempfile::tempdir().unwrap().into_path(); println!("Testing image extraction from: {}", test_pdf_path.display()); // Now try image extraction let result = pdf_tool( test_pdf_path.to_str().unwrap(), "extract_images", &cache_dir, ) .await; println!("Image extraction result: {:?}", result); assert!(result.is_ok(), "PDF image extraction should succeed"); let content = result.unwrap(); assert!( !content.is_empty(), "Image extraction result should not be empty" ); let text = content[0].as_text().unwrap(); println!("Extracted content: {}", text); // Should either find images or explicitly state none were found assert!( text.contains("Saved image to:") || text.contains("No images found"), "Should either save images or report none found" ); // If we found images, verify they exist if text.contains("Saved image to:") { // Extract the file path from the output let file_path = text .lines() .find(|line| line.contains("Saved image to:")) .and_then(|line| line.split(": ").nth(1)) .and_then(|path| path.split(" (").next()) .expect("Should have a valid file path"); println!("Verifying image file exists: {}", file_path); assert!(PathBuf::from(file_path).exists(), "Image file should exist"); } } #[tokio::test] async fn test_pdf_invalid_path() { let cache_dir = tempfile::tempdir().unwrap().into_path(); let result = pdf_tool("nonexistent.pdf", "extract_text", &cache_dir).await; assert!(result.is_err(), "Should fail with invalid path"); } #[tokio::test] async fn test_pdf_invalid_operation() { let test_pdf_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("src/computercontroller/tests/data/test.pdf"); let cache_dir = tempfile::tempdir().unwrap().into_path(); let result = pdf_tool( test_pdf_path.to_str().unwrap(), "invalid_operation", &cache_dir, ) .await; assert!(result.is_err(), "Should fail with invalid operation"); } }