use std::io::{self, Write};
use std::path::Path;
use std::process::ExitCode;
use docling::{DocumentConverter, ImageMode, InputFormat, Pipeline, SourceDocument};
fn main() -> ExitCode {
{
let mut args = std::env::args().skip(1);
if args.next().as_deref() == Some("serve") {
return run_serve(args.collect());
}
}
let mut strict = false;
let mut to = "md".to_string();
let mut images = "placeholder".to_string();
let mut fetch_images = false;
let mut no_stream = false;
let mut no_table_former = false;
let mut no_ocr = false;
let mut use_web_browser = false;
let mut enrich_picture_classes = false;
let mut enrich_code = false;
let mut enrich_formula = false;
let mut bench_warm: Option<usize> = None;
let mut path: Option<String> = None;
let mut args = std::env::args().skip(1);
while let Some(arg) = args.next() {
match arg.as_str() {
"--strict" => strict = true,
"--fetch-images" => fetch_images = true,
"--no-stream" => no_stream = true,
"--no-table-former" => no_table_former = true,
"--no-ocr" => no_ocr = true,
"--use-web-browser" => use_web_browser = true,
"--enrich-picture-classes" => enrich_picture_classes = true,
"--enrich-code" => enrich_code = true,
"--enrich-formula" => enrich_formula = true,
"--to" => to = args.next().unwrap_or_default(),
"--images" => images = args.next().unwrap_or_default(),
"--bench-warm" => {
bench_warm = args.next().and_then(|n| n.parse::<usize>().ok());
if bench_warm.is_none() {
eprintln!("error: --bench-warm needs a positive run count");
return ExitCode::from(2);
}
}
_ => path = Some(arg),
}
}
if !matches!(to.as_str(), "md" | "markdown" | "json" | "dclx" | "chunks") {
eprintln!("error: unknown --to '{to}' (expected: md, json, dclx, chunks)");
return ExitCode::from(2);
}
let image_mode = match images.as_str() {
"placeholder" => ImageMode::Placeholder,
"embedded" => ImageMode::Embedded,
"referenced" => ImageMode::Referenced,
other => {
eprintln!(
"error: unknown --images '{other}' (expected: placeholder, embedded, referenced)"
);
return ExitCode::from(2);
}
};
let Some(path) = path else {
eprintln!("usage: docling-rs [--strict] [--to md|json] [--images MODE] [--fetch-images] [--no-stream] [--no-table-former] [--no-ocr] [--use-web-browser] <input-file>");
return ExitCode::from(2);
};
let source = match SourceDocument::from_file(&path) {
Ok(src) => src,
Err(e) => {
eprintln!("error: {e}");
return ExitCode::FAILURE;
}
};
if let Some(runs) = bench_warm {
return match bench_warm_conversion(&source, runs, no_table_former, no_ocr) {
Ok(avg) => {
println!("{avg:.6}");
eprintln!(
"warm conversion: {:.4}s/doc over {runs} runs (startup excluded)",
avg
);
ExitCode::SUCCESS
}
Err(e) => {
eprintln!("error: {e}");
ExitCode::FAILURE
}
};
}
let converter = DocumentConverter::new()
.strict(strict)
.fetch_images(fetch_images)
.no_table_former(no_table_former)
.no_ocr(no_ocr)
.use_web_browser(use_web_browser)
.do_picture_classification(enrich_picture_classes)
.do_code_enrichment(enrich_code)
.do_formula_enrichment(enrich_formula);
let is_markdown = matches!(to.as_str(), "md" | "markdown");
if is_markdown && image_mode != ImageMode::Referenced && !no_stream {
let stream = match converter.convert_streaming_images(source, image_mode) {
Ok(s) => s,
Err(e) => {
eprintln!("error: {e}");
return ExitCode::FAILURE;
}
};
let stdout = io::stdout();
let mut out = io::BufWriter::new(stdout.lock());
for chunk in stream {
match chunk {
Ok(s) => {
if let Err(e) = out.write_all(s.as_bytes()) {
eprintln!("error: writing output: {e}");
return ExitCode::FAILURE;
}
}
Err(e) => {
let _ = out.flush();
eprintln!("error: {e}");
return ExitCode::FAILURE;
}
}
}
if let Err(e) = out.flush() {
eprintln!("error: writing output: {e}");
return ExitCode::FAILURE;
}
return ExitCode::SUCCESS;
}
let document = match converter.convert(source) {
Ok(result) => result.document,
Err(e) => {
eprintln!("error: {e}");
return ExitCode::FAILURE;
}
};
if to == "json" {
println!("{}", document.export_to_json());
return ExitCode::SUCCESS;
}
if to == "chunks" {
print!("{}", chunks_json(&document));
return ExitCode::SUCCESS;
}
if to == "dclx" {
let stem = Path::new(&path)
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| "document".into());
let out = std::path::PathBuf::from(format!("{stem}.dclx"));
if let Err(e) = docling::dclx::save_as_dclx(&document, &out) {
eprintln!("error: dclx: {e}");
return ExitCode::FAILURE;
}
println!("{}", out.display());
return ExitCode::SUCCESS;
}
if image_mode == ImageMode::Placeholder {
print!("{}", document.export_to_markdown());
return ExitCode::SUCCESS;
}
let (md, artifacts) = document.export_to_markdown_with_images(image_mode, "artifacts");
for (rel, bytes) in &artifacts {
let rel = Path::new(rel);
if let Some(dir) = rel.parent() {
if let Err(e) = std::fs::create_dir_all(dir) {
eprintln!("error: creating {}: {e}", dir.display());
return ExitCode::FAILURE;
}
}
if let Err(e) = std::fs::write(rel, bytes) {
eprintln!("error: writing {}: {e}", rel.display());
return ExitCode::FAILURE;
}
}
if !artifacts.is_empty() {
eprintln!("wrote {} image(s) to ./artifacts/", artifacts.len());
}
print!("{md}");
ExitCode::SUCCESS
}
#[cfg(feature = "serve")]
fn run_serve(args: Vec<String>) -> ExitCode {
use docling_serve::ServeConfig;
let mut cfg = ServeConfig::default();
let mut it = args.into_iter();
while let Some(arg) = it.next() {
match arg.as_str() {
"--addr" => match it.next() {
Some(v) => cfg.addr = v,
None => return serve_usage("--addr needs HOST:PORT"),
},
"--concurrency" => match it.next().and_then(|v| v.parse().ok()) {
Some(v) if v >= 1 => cfg.concurrency = v,
_ => return serve_usage("--concurrency needs a positive integer"),
},
"--max-body-mb" => match it.next().and_then(|v| v.parse::<usize>().ok()) {
Some(v) if v >= 1 => cfg.max_body_bytes = v * 1024 * 1024,
_ => return serve_usage("--max-body-mb needs a positive integer"),
},
"--warmup" => cfg.warmup = true,
"--no-url-fetch" => cfg.allow_url_fetch = false,
"--strict" => cfg.strict = true,
other => return serve_usage(&format!("unknown argument '{other}'")),
}
}
let runtime = match tokio::runtime::Runtime::new() {
Ok(rt) => rt,
Err(e) => {
eprintln!("error: tokio runtime: {e}");
return ExitCode::FAILURE;
}
};
match runtime.block_on(docling_serve::serve(cfg)) {
Ok(()) => ExitCode::SUCCESS,
Err(e) => {
eprintln!("error: {e}");
ExitCode::FAILURE
}
}
}
#[cfg(feature = "serve")]
fn serve_usage(err: &str) -> ExitCode {
eprintln!("error: {err}");
eprintln!("usage: docling-rs serve [--addr HOST:PORT] [--concurrency N] [--max-body-mb N] [--warmup] [--no-url-fetch] [--strict]");
ExitCode::from(2)
}
#[cfg(not(feature = "serve"))]
fn run_serve(_args: Vec<String>) -> ExitCode {
eprintln!(
"error: this binary was built without the HTTP server.\n\
Rebuild with `cargo build -p docling-cli --features serve`, or use the\n\
standalone server: `cargo run -p docling-serve --release -- --help`."
);
ExitCode::from(2)
}
fn bench_warm_conversion(
source: &SourceDocument,
runs: usize,
no_table_former: bool,
no_ocr: bool,
) -> Result<f64, String> {
let mut pipeline = Pipeline::new()
.map_err(|e| e.to_string())?
.no_table_former(no_table_former)
.no_ocr(no_ocr);
let once = |p: &mut Pipeline| -> Result<(), String> {
match source.format {
InputFormat::Pdf => p
.convert(&source.bytes, None, &source.name)
.map(|_| ())
.map_err(|e| e.to_string()),
InputFormat::Image => p
.convert_image(&source.bytes, &source.name)
.map(|_| ())
.map_err(|e| e.to_string()),
other => Err(format!(
"--bench-warm supports PDF/image only, not {other:?}"
)),
}
};
once(&mut pipeline)?; let mut total = 0.0f64;
for _ in 0..runs {
let t = std::time::Instant::now();
once(&mut pipeline)?;
total += t.elapsed().as_secs_f64();
}
Ok(total / runs as f64)
}
fn chunks_json(document: &docling::DoclingDocument) -> String {
let mut warn = |msg: String| eprintln!("warning: {msg}");
let out = docling::chunks::chunk_records(document, &mut warn);
format!(
"{}\n",
serde_json::to_string_pretty(&out).expect("chunks are serializable")
)
}