use crate::DriverCallback;
use crate::DriverContext;
use crate::{
DriverCategory,
types::{Driver, DriverParameter},
};
use crate::{ensure_dir, file_exists, validate_path};
use anyhow::Result;
use serde_json::{Value, json};
use std::collections::HashMap;
#[derive(Debug)]
pub struct PdfReadDriver;
#[async_trait::async_trait]
impl Driver for PdfReadDriver {
fn name(&self) -> &str {
"pdf_read"
}
fn description(&self) -> &str {
"Read and extract text content from PDF files"
}
fn usage_hint(&self) -> &str {
"Use this skill when the user wants to read PDF documents, extract text from PDF files"
}
fn parameters(&self) -> Vec<DriverParameter> {
vec![
DriverParameter {
name: "path".to_string(),
param_type: "string".to_string(),
description: "Path to the PDF file".to_string(),
required: true,
default: None,
example: Some(Value::String("document.pdf".to_string())),
enum_values: None,
},
DriverParameter {
name: "start_page".to_string(),
param_type: "integer".to_string(),
description: "Starting page number (1-indexed)".to_string(),
required: false,
default: Some(Value::Number(1.into())),
example: Some(Value::Number(1.into())),
enum_values: None,
},
DriverParameter {
name: "end_page".to_string(),
param_type: "integer".to_string(),
description: "Ending page number (inclusive)".to_string(),
required: false,
default: None,
example: Some(Value::Number(10.into())),
enum_values: None,
},
]
}
fn example_call(&self) -> Value {
json!({
"action": "pdf_read",
"parameters": {
"path": "document.pdf"
}
})
}
fn example_output(&self) -> String {
"PDF content extracted from document.pdf\nPage 1: This is the content...".to_string()
}
fn category(&self) -> DriverCategory {
DriverCategory::Document
}
async fn execute(
&self,
parameters: &HashMap<String, Value>,
callback: Option<&dyn DriverCallback>,
context: Option<&DriverContext>,
) -> Result<String> {
let task_id = context.as_ref().and_then(|c| c.task_id()).map(String::from);
let driver_index = context.as_ref().and_then(|c| c.driver_index());
let step_name = context
.as_ref()
.and_then(|c| c.driver_name())
.map(String::from);
let cb = callback;
if let Some(cb) = cb {
cb.on_start(task_id.clone(), driver_index, step_name);
cb.on_log(
task_id.clone(),
driver_index,
Some("Starting PDF read operation".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(10), None);
}
let path = parameters
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Reading PDF file: {}", path)),
);
cb.on_progress(task_id.clone(), driver_index, Some(20), None);
}
let start_page = parameters
.get("start_page")
.and_then(|v| v.as_u64())
.unwrap_or(1) as usize;
let end_page = parameters
.get("end_page")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!(
"Start page: {}, End page: {:?}",
start_page, end_page
)),
);
cb.on_progress(task_id.clone(), driver_index, Some(30), None);
}
let validated_path = validate_path(path, None)?;
if !file_exists(&validated_path.to_string_lossy()) {
anyhow::bail!("PDF file not found: {}", path);
}
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("Validated path, loading PDF".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(40), None);
}
use pdf_extract::extract_text;
let full_text = extract_text(&validated_path)
.map_err(|e| anyhow::anyhow!("Failed to extract PDF text: {}", e))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("PDF text extracted successfully".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(60), None);
}
let pages: Vec<&str> = full_text.split("\n\n").collect();
let start = start_page.saturating_sub(1);
let end = end_page.unwrap_or(pages.len()).min(pages.len());
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!(
"Total pages: {}, Selected range: {}-{}",
pages.len(),
start + 1,
end
)),
);
cb.on_progress(task_id.clone(), driver_index, Some(75), None);
}
if start >= pages.len() {
anyhow::bail!(
"Start page {} exceeds total pages {}",
start_page,
pages.len()
);
}
let mut output = format!("PDF file: {}\n", path);
output.push_str(&format!("Total pages: {}\n", pages.len()));
output.push_str(&format!("Showing pages {}-{}\n\n", start + 1, end));
for i in start..end {
output.push_str(&format!("=== Page {} ===\n", i + 1));
output.push_str(pages[i]);
output.push_str("\n\n");
}
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!(
"Completed PDF read, output length: {} bytes",
output.len()
)),
);
cb.on_progress(task_id.clone(), driver_index, Some(100), None);
cb.on_complete(
task_id.clone(),
driver_index,
Some("pdf_read".to_string()),
Some(output.clone()),
);
}
Ok(output)
}
fn validate(&self, parameters: &HashMap<String, Value>) -> Result<()> {
parameters
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing required parameter: path"))?;
Ok(())
}
}
#[derive(Debug)]
pub struct PdfMergeDriver;
#[async_trait::async_trait]
impl Driver for PdfMergeDriver {
fn name(&self) -> &str {
"pdf_merge"
}
fn description(&self) -> &str {
"Merge multiple PDF files into a single PDF"
}
fn usage_hint(&self) -> &str {
"Use this skill when the user wants to combine multiple PDF files into one"
}
fn parameters(&self) -> Vec<DriverParameter> {
vec![
DriverParameter {
name: "inputs".to_string(),
param_type: "array".to_string(),
description: "Array of PDF file paths to merge".to_string(),
required: true,
default: None,
example: Some(json!(["file1.pdf", "file2.pdf", "file3.pdf"])),
enum_values: None,
},
DriverParameter {
name: "output".to_string(),
param_type: "string".to_string(),
description: "Output PDF file path".to_string(),
required: true,
default: None,
example: Some(Value::String("merged.pdf".to_string())),
enum_values: None,
},
]
}
fn example_call(&self) -> Value {
json!({
"action": "pdf_merge",
"parameters": {
"inputs": ["doc1.pdf", "doc2.pdf"],
"output": "merged.pdf"
}
})
}
fn example_output(&self) -> String {
"Merged 2 PDF files into: merged.pdf".to_string()
}
fn category(&self) -> DriverCategory {
DriverCategory::Document
}
async fn execute(
&self,
parameters: &HashMap<String, Value>,
callback: Option<&dyn DriverCallback>,
context: Option<&DriverContext>,
) -> Result<String> {
let task_id = context.as_ref().and_then(|c| c.task_id()).map(String::from);
let driver_index = context.as_ref().and_then(|c| c.driver_index());
let step_name = context
.as_ref()
.and_then(|c| c.driver_name())
.map(String::from);
let cb = callback;
if let Some(cb) = cb {
cb.on_start(task_id.clone(), driver_index, step_name);
cb.on_log(
task_id.clone(),
driver_index,
Some("Starting PDF merge operation".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(10), None);
}
use lopdf::{Document, Object, ObjectId};
let inputs = parameters
.get("inputs")
.ok_or_else(|| anyhow::anyhow!("Missing 'inputs' parameter"))?
.as_array()
.ok_or_else(|| anyhow::anyhow!("'inputs' must be an array"))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Number of input files: {}", inputs.len())),
);
cb.on_progress(task_id.clone(), driver_index, Some(15), None);
}
let output = parameters
.get("output")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing 'output' parameter"))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Output file: {}", output)),
);
cb.on_progress(task_id.clone(), driver_index, Some(20), None);
}
if inputs.is_empty() {
anyhow::bail!("At least one input file is required");
}
let validated_output = validate_path(output, None)?;
if let Some(parent) = validated_output.parent() {
ensure_dir(&parent.to_string_lossy())?;
}
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("Validated output path, creating directory".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(25), None);
}
let mut merged_doc = Document::new();
let mut total_pages = 0;
let mut max_id = 0;
let total_inputs = inputs.len();
for (idx, input_path) in inputs.iter().enumerate() {
let path = input_path
.as_str()
.ok_or_else(|| anyhow::anyhow!("Input path must be a string"))?;
if let Some(cb) = cb {
let progress = 25 + ((idx + 1) as f32 / total_inputs as f32 * 50.0) as u32;
cb.on_log(
task_id.clone(),
driver_index,
Some(format!(
"Processing file {}/{}: {}",
idx + 1,
total_inputs,
path
)),
);
cb.on_progress(task_id.clone(), driver_index, Some(progress), None);
}
let validated_input = validate_path(path, None)?;
let doc = Document::load(&validated_input)
.map_err(|e| anyhow::anyhow!("Failed to load PDF '{}': {}", path, e))?;
let pages = doc.page_iter().collect::<Vec<_>>();
total_pages += pages.len();
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Loaded PDF with {} pages", pages.len())),
);
}
for (id, object) in doc.objects.iter() {
let new_id = (id.0 + max_id, id.1 + max_id as u16);
merged_doc.objects.insert(new_id, object.clone());
}
max_id += doc.max_id;
}
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("All files loaded, total pages: {}", total_pages)),
);
cb.on_progress(task_id.clone(), driver_index, Some(80), None);
}
let mut page_objects = Vec::new();
for (object_id, object) in merged_doc.objects.iter() {
if let Ok(dict) = object.as_dict() {
if let Ok(value) = dict.get(b"Type") {
if let Ok(name) = value.as_name() {
if name == b"Page" {
page_objects.push(*object_id);
}
}
}
}
}
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Found {} page objects", page_objects.len())),
);
cb.on_progress(task_id.clone(), driver_index, Some(90), None);
}
if page_objects.is_empty() {
anyhow::bail!("No pages found in input PDFs");
}
merged_doc
.save(&validated_output)
.map_err(|e| anyhow::anyhow!("Failed to save merged PDF: {}", e))?;
let result = format!(
"Merged {} PDF files into: {} ({} total pages)",
inputs.len(),
output,
total_pages
);
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Merge completed: {}", result)),
);
cb.on_progress(task_id.clone(), driver_index, Some(100), None);
cb.on_complete(
task_id.clone(),
driver_index,
Some("pdf_merge".to_string()),
Some(result.clone()),
);
}
Ok(result)
}
fn validate(&self, parameters: &HashMap<String, Value>) -> Result<()> {
parameters
.get("inputs")
.ok_or_else(|| anyhow::anyhow!("Missing required parameter: inputs"))?;
parameters
.get("output")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing required parameter: output"))?;
Ok(())
}
}
#[derive(Debug)]
pub struct PdfInfoDriver;
#[async_trait::async_trait]
impl Driver for PdfInfoDriver {
fn name(&self) -> &str {
"pdf_info"
}
fn description(&self) -> &str {
"Get metadata information from PDF file (pages, title, author, etc.)"
}
fn usage_hint(&self) -> &str {
"Use this skill when the user wants to get information about a PDF file"
}
fn parameters(&self) -> Vec<DriverParameter> {
vec![DriverParameter {
name: "path".to_string(),
param_type: "string".to_string(),
description: "Path to the PDF file".to_string(),
required: true,
default: None,
example: Some(Value::String("document.pdf".to_string())),
enum_values: None,
}]
}
fn example_call(&self) -> Value {
json!({
"action": "pdf_info",
"parameters": {
"path": "document.pdf"
}
})
}
fn example_output(&self) -> String {
"PDF Info:\nPages: 25\nTitle: My Document\nAuthor: John Doe".to_string()
}
fn category(&self) -> DriverCategory {
DriverCategory::Document
}
async fn execute(
&self,
parameters: &HashMap<String, Value>,
callback: Option<&dyn DriverCallback>,
context: Option<&DriverContext>,
) -> Result<String> {
let task_id = context.as_ref().and_then(|c| c.task_id()).map(String::from);
let driver_index = context.as_ref().and_then(|c| c.driver_index());
let step_name = context
.as_ref()
.and_then(|c| c.driver_name())
.map(String::from);
let cb = callback;
if let Some(cb) = cb {
cb.on_start(task_id.clone(), driver_index, step_name);
cb.on_log(
task_id.clone(),
driver_index,
Some("Starting PDF info operation".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(10), None);
}
use lopdf::Document;
let path = parameters
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing 'path' parameter"))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Reading PDF info for: {}", path)),
);
cb.on_progress(task_id.clone(), driver_index, Some(20), None);
}
let validated_path = validate_path(path, None)?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("Validated path, loading PDF".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(30), None);
}
let doc = Document::load(&validated_path)
.map_err(|e| anyhow::anyhow!("Failed to load PDF: {}", e))?;
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("PDF loaded successfully".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(50), None);
}
let pages = doc.page_iter().count();
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!("Total pages: {}", pages)),
);
cb.on_progress(task_id.clone(), driver_index, Some(60), None);
}
let mut output = format!("PDF Info for: {}\n", path);
output.push_str(&format!("Total pages: {}\n", pages));
if let Ok(info_ref) = doc.trailer.get(b"Info") {
if let Ok(info_id) = info_ref.as_reference() {
if let Ok(info) = doc.get_object(info_id) {
if let Ok(dict) = info.as_dict() {
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some("Extracting metadata from PDF".to_string()),
);
cb.on_progress(task_id.clone(), driver_index, Some(70), None);
}
if let Ok(title) = dict.get(b"Title") {
if let Ok(title_str) = title.as_str() {
output.push_str(&format!("Title: {:?}\n", title_str));
}
}
if let Ok(author) = dict.get(b"Author") {
if let Ok(author_str) = author.as_str() {
output.push_str(&format!("Author: {:?}\n", author_str));
}
}
if let Ok(subject) = dict.get(b"Subject") {
if let Ok(subject_str) = subject.as_str() {
output.push_str(&format!("Subject: {:?}\n", subject_str));
}
}
if let Ok(creator) = dict.get(b"Creator") {
if let Ok(creator_str) = creator.as_str() {
output.push_str(&format!("Creator: {:?}\n", creator_str));
}
}
if let Ok(producer) = dict.get(b"Producer") {
if let Ok(producer_str) = producer.as_str() {
output.push_str(&format!("Producer: {:?}\n", producer_str));
}
}
}
}
}
}
let file_size = std::fs::metadata(&validated_path)
.map(|m| m.len())
.unwrap_or(0);
output.push_str(&format!("File size: {:.2} KB\n", file_size as f64 / 1024.0));
if let Some(cb) = cb {
cb.on_log(
task_id.clone(),
driver_index,
Some(format!(
"PDF info completed, output length: {} bytes",
output.len()
)),
);
cb.on_progress(task_id.clone(), driver_index, Some(100), None);
cb.on_complete(
task_id.clone(),
driver_index,
Some("pdf_info".to_string()),
Some(output.clone()),
);
}
Ok(output)
}
fn validate(&self, parameters: &HashMap<String, Value>) -> Result<()> {
parameters
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("Missing required parameter: path"))?;
Ok(())
}
}