datalab-cli 0.1.0

A powerful CLI for converting, extracting, and processing documents using the Datalab API
Documentation
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 {
    /// File path or URL to convert
    #[arg(value_name = "FILE|URL")]
    pub input: String,

    /// Output format: markdown, html, json, chunks
    #[arg(
        long,
        default_value = "markdown",
        value_name = "FORMAT",
        help_heading = "Output Options"
    )]
    pub output_format: String,

    /// Processing mode: fast, balanced, accurate
    #[arg(
        long,
        default_value = "fast",
        value_name = "MODE",
        help_heading = "Processing Options"
    )]
    pub mode: String,

    /// Maximum pages to process
    #[arg(long, value_name = "N", help_heading = "Processing Options")]
    pub max_pages: Option<u32>,

    /// Page range (e.g., "0-5,10")
    #[arg(long, value_name = "RANGE", help_heading = "Processing Options")]
    pub page_range: Option<String>,

    /// Add page delimiters to output
    #[arg(long, help_heading = "Output Options")]
    pub paginate: bool,

    /// Skip local cache lookup
    #[arg(long, help_heading = "Cache Options")]
    pub skip_cache: bool,

    /// Force reprocessing (skip API cache)
    #[arg(long, help_heading = "Cache Options")]
    pub force: bool,

    /// Save checkpoint for follow-up extraction/segmentation
    #[arg(long, help_heading = "Advanced Options")]
    pub save_checkpoint: bool,

    /// Extra features: track_changes, chart_understanding, extract_links
    #[arg(long, value_name = "FEATURES", help_heading = "Advanced Options")]
    pub extras: Option<String>,

    /// Add block IDs for citation tracking
    #[arg(long, help_heading = "Advanced Options")]
    pub add_block_ids: bool,

    /// Use token-efficient markdown format
    #[arg(long, help_heading = "Output Options")]
    pub token_efficient_markdown: bool,

    /// Disable image extraction from document
    #[arg(long, help_heading = "Advanced Options")]
    pub disable_image_extraction: bool,

    /// Disable image caption generation
    #[arg(long, help_heading = "Advanced Options")]
    pub disable_image_captions: bool,

    /// Write result to file instead of stdout
    #[arg(long, short, value_name = "FILE", help_heading = "Output Options")]
    pub output: Option<PathBuf>,

    /// Request timeout in seconds
    #[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))
    };

    // Emit start progress
    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(())
}