use crate::cache::Cache;
use crate::client::{add_form_field, build_form_with_file, DatalabClient};
use crate::error::{DatalabError, Result};
use crate::output::Progress;
use clap::Args;
use reqwest::multipart::Form;
use serde_json::json;
use std::fs;
use std::path::PathBuf;
#[derive(Args, Debug)]
#[command(after_help = "EXAMPLES:\n \
# Basic conversion to markdown\n \
datalab convert document.pdf\n\n \
# High-quality conversion with charts\n \
datalab convert report.pdf --mode accurate --extras chart_understanding\n\n \
# Convert specific pages\n \
datalab convert book.pdf --page-range \"0-10,50-60\" --paginate")]
pub struct ConvertArgs {
#[arg(value_name = "FILE|URL")]
pub input: String,
#[arg(
long,
default_value = "markdown",
value_name = "FORMAT",
help_heading = "Output Options"
)]
pub output_format: String,
#[arg(
long,
default_value = "fast",
value_name = "MODE",
help_heading = "Processing Options"
)]
pub mode: String,
#[arg(long, value_name = "N", help_heading = "Processing Options")]
pub max_pages: Option<u32>,
#[arg(long, value_name = "RANGE", help_heading = "Processing Options")]
pub page_range: Option<String>,
#[arg(long, help_heading = "Output Options")]
pub paginate: bool,
#[arg(long, help_heading = "Cache Options")]
pub skip_cache: bool,
#[arg(long, help_heading = "Cache Options")]
pub force: bool,
#[arg(long, help_heading = "Advanced Options")]
pub save_checkpoint: bool,
#[arg(long, value_name = "FEATURES", help_heading = "Advanced Options")]
pub extras: Option<String>,
#[arg(long, help_heading = "Advanced Options")]
pub add_block_ids: bool,
#[arg(long, help_heading = "Output Options")]
pub token_efficient_markdown: bool,
#[arg(long, help_heading = "Advanced Options")]
pub disable_image_extraction: bool,
#[arg(long, help_heading = "Advanced Options")]
pub disable_image_captions: bool,
#[arg(long, short, value_name = "FILE", help_heading = "Output Options")]
pub output: Option<PathBuf>,
#[arg(
long,
default_value = "300",
value_name = "SECS",
help_heading = "Advanced Options"
)]
pub timeout: u64,
}
impl ConvertArgs {
fn to_cache_params(&self) -> serde_json::Value {
json!({
"output_format": self.output_format,
"mode": self.mode,
"max_pages": self.max_pages,
"page_range": self.page_range,
"paginate": self.paginate,
"save_checkpoint": self.save_checkpoint,
"extras": self.extras,
"add_block_ids": self.add_block_ids,
"token_efficient_markdown": self.token_efficient_markdown,
"disable_image_extraction": self.disable_image_extraction,
"disable_image_captions": self.disable_image_captions,
})
}
fn add_to_form(&self, mut form: Form) -> Form {
form = add_form_field(form, "output_format", &self.output_format);
form = add_form_field(form, "mode", &self.mode);
if let Some(max_pages) = self.max_pages {
form = add_form_field(form, "max_pages", &max_pages.to_string());
}
if let Some(ref page_range) = self.page_range {
form = add_form_field(form, "page_range", page_range);
}
if self.paginate {
form = add_form_field(form, "paginate", "true");
}
if self.force {
form = add_form_field(form, "skip_cache", "true");
}
if self.save_checkpoint {
form = add_form_field(form, "save_checkpoint", "true");
}
if let Some(ref extras) = self.extras {
form = add_form_field(form, "extras", extras);
}
if self.add_block_ids {
form = add_form_field(form, "add_block_ids", "true");
}
if self.token_efficient_markdown {
form = add_form_field(form, "token_efficient_markdown", "true");
}
if self.disable_image_extraction {
form = add_form_field(form, "disable_image_extraction", "true");
}
if self.disable_image_captions {
form = add_form_field(form, "disable_image_captions", "true");
}
form
}
}
pub async fn execute(args: ConvertArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
let cache = Cache::new()?;
let is_url = args.input.starts_with("http://") || args.input.starts_with("https://");
let file_path = if is_url {
None
} else {
Some(PathBuf::from(&args.input))
};
let file_str = file_path.as_ref().map(|p| p.to_string_lossy().to_string());
progress.start("convert", file_str.as_deref());
let file_hash = if let Some(ref path) = file_path {
if !path.exists() {
return Err(DatalabError::FileNotFound(path.clone()));
}
Some(Cache::hash_file(path)?)
} else {
None
};
let cache_params = args.to_cache_params();
let cache_key = Cache::generate_key(
file_hash.as_deref(),
if is_url { Some(&args.input) } else { None },
"convert",
&cache_params,
);
if !args.skip_cache {
if let Some(cached) = cache.get(&cache_key) {
progress.cache_hit(&cache_key);
output_result(&cached, args.output.as_ref())?;
return Ok(());
}
}
let form = if let Some(ref path) = file_path {
let (form, _) = build_form_with_file(path)?;
args.add_to_form(form)
} else {
let form = Form::new().text("file_url", args.input.clone());
args.add_to_form(form)
};
let result = client.submit_and_poll("convert", form, progress).await?;
let file_path_str = file_path.as_ref().map(|p| p.to_string_lossy().to_string());
cache.set(
&cache_key,
&result,
"convert",
file_hash.as_deref(),
file_path_str.as_deref(),
)?;
output_result(&result, args.output.as_ref())?;
Ok(())
}
fn output_result(result: &serde_json::Value, output_file: Option<&PathBuf>) -> Result<()> {
let json_output = serde_json::to_string_pretty(result)?;
if let Some(path) = output_file {
fs::write(path, &json_output)?;
} else {
println!("{}", json_output);
}
Ok(())
}