cloudiful_docling_convert/processor/service/

output.rs

use std::path::{Path, PathBuf};

use serde_json::Value;
use tokio::fs;

use super::DocumentConverter;
use crate::api::DoclingResult;
use crate::document::{
    ConvertedDocument, ConvertedDocumentMetadata, InputDocument, InputKind, OutputFormat,
};
use crate::error::{PdfConvertError, Result};
use crate::models::{ChunkDocumentResponse, ExportDocumentResponse};

impl DocumentConverter {
    pub fn calculate_output_path(
        output_dir: &Path,
        filename: &str,
        output_format: OutputFormat,
    ) -> PathBuf {
        let stem = Path::new(filename)
            .file_stem()
            .and_then(|value| value.to_str())
            .unwrap_or("output");
        output_dir.join(format!("{}.{}", stem, output_format.extension()))
    }

    pub(super) async fn write_output_file(
        output_path: &Path,
        document: &ConvertedDocument,
        output_format: OutputFormat,
    ) -> Result<()> {
        if let Some(parent) = output_path.parent() {
            fs::create_dir_all(parent).await.map_err(|error| {
                PdfConvertError::io_error(
                    format!("creating output directory: {}", parent.display()),
                    error,
                )
            })?;
        }

        let content = select_output_content(document, output_format)?;
        fs::write(output_path, content).await.map_err(|error| {
            PdfConvertError::io_error(
                format!("writing output file: {}", output_path.display()),
                error,
            )
        })?;
        Ok(())
    }

    pub(super) fn build_document(
        input: &InputDocument,
        input_kind: InputKind,
        result: DoclingResult,
    ) -> ConvertedDocument {
        let mut document = empty_document(input, input_kind);
        match result {
            DoclingResult::Convert(response) => {
                document.filename = response.document.filename.clone();
                apply_export_document(&mut document, response.document);
                document.errors = response
                    .errors
                    .iter()
                    .filter_map(|error| error.error_message.clone())
                    .collect();
            }
            DoclingResult::Chunk(response) => {
                apply_chunk_response(&mut document, response);
            }
            DoclingResult::Json(value) => {
                apply_json_result(&mut document, value);
            }
            DoclingResult::Zip(bytes) => {
                document.archive = Some(bytes.to_vec());
            }
            DoclingResult::Failure(failure) => {
                document.errors.push(failure.failure.message);
            }
        }
        document
    }
}

fn empty_document(input: &InputDocument, input_kind: InputKind) -> ConvertedDocument {
    ConvertedDocument {
        filename: input.filename.clone(),
        markdown: None,
        text: None,
        json: None,
        html: None,
        doctags: None,
        doclang: None,
        chunks: Vec::new(),
        chunk_response: None,
        archive: None,
        metadata: ConvertedDocumentMetadata {
            input_kind,
            media_type: input.media_type.clone(),
        },
        errors: Vec::new(),
    }
}

fn apply_export_document(document: &mut ConvertedDocument, response: ExportDocumentResponse) {
    document.markdown = response.md_content.clone();
    document.text = response.text_content.clone();
    document.json = response.json_content.clone();
    document.html = response.html_content.clone();
    document.doctags = response.doctags_content.clone();
    document.doclang = response.doclang_content.clone();
}

fn apply_chunk_response(document: &mut ConvertedDocument, response: ChunkDocumentResponse) {
    document.filename = response
        .chunks
        .first()
        .map(|chunk| chunk.filename.clone())
        .or_else(|| {
            response
                .documents
                .first()
                .map(|item| item.document.filename.clone())
        })
        .unwrap_or_else(|| document.filename.clone());
    document.chunks = response.chunks.clone();
    document.errors = response
        .documents
        .iter()
        .flat_map(|item| item.errors.iter())
        .filter_map(|error| error.error_message.clone())
        .collect();
    if let Some(item) = response.documents.first().cloned() {
        apply_export_document(document, item.document);
    }
    document.chunk_response = Some(response);
}

fn apply_json_result(document: &mut ConvertedDocument, value: Value) {
    let source = value.get("document").unwrap_or(&value);
    if let Some(object) = source.as_object() {
        document.markdown = object
            .get("md_content")
            .and_then(Value::as_str)
            .map(ToString::to_string);
        document.text = object
            .get("text_content")
            .and_then(Value::as_str)
            .map(ToString::to_string);
        document.html = object
            .get("html_content")
            .and_then(Value::as_str)
            .map(ToString::to_string);
        document.doctags = object
            .get("doctags_content")
            .and_then(Value::as_str)
            .map(ToString::to_string);
        document.doclang = object
            .get("doclang_content")
            .and_then(Value::as_str)
            .map(ToString::to_string);
        document.json = object.get("json_content").cloned();
    }
    document.errors = collect_json_errors(&value);
}

fn collect_json_errors(value: &Value) -> Vec<String> {
    let mut errors = Vec::new();
    if let Some(message) = value.get("error_message").and_then(Value::as_str) {
        errors.push(message.to_string());
    }
    if let Some(items) = value.get("errors").and_then(Value::as_array) {
        errors.extend(items.iter().filter_map(|item| {
            item.get("error_message")
                .or_else(|| item.get("message"))
                .and_then(Value::as_str)
                .map(ToString::to_string)
        }));
    }
    errors
}

fn select_output_content(
    document: &ConvertedDocument,
    output_format: OutputFormat,
) -> Result<Vec<u8>> {
    if output_format.is_archive() {
        return document.archive.clone().ok_or_else(|| {
            PdfConvertError::operation_error(
                "writing archive",
                "Docling did not return an application/zip response",
            )
        });
    }

    if output_format.is_chunk_output() {
        return document
            .chunk_response
            .as_ref()
            .map(serde_json::to_vec_pretty)
            .transpose()
            .map_err(PdfConvertError::from)?
            .ok_or_else(|| {
                PdfConvertError::operation_error(
                    "writing chunks",
                    "Docling did not return a chunk response",
                )
            });
    }

    match output_format {
        OutputFormat::Md => text_content(document.markdown.as_deref(), "markdown"),
        OutputFormat::Text => text_content(document.text.as_deref(), "text"),
        OutputFormat::Html => text_content(document.html.as_deref(), "html"),
        OutputFormat::Doctags => text_content(document.doctags.as_deref(), "doctags"),
        OutputFormat::Doclang => text_content(document.doclang.as_deref(), "doclang"),
        OutputFormat::Json => document
            .json
            .as_ref()
            .map(serde_json::to_vec_pretty)
            .transpose()
            .map_err(PdfConvertError::from)?
            .ok_or_else(|| {
                PdfConvertError::operation_error("writing json", "JSON output is empty")
            }),
        OutputFormat::Yaml
        | OutputFormat::HtmlSplitPage
        | OutputFormat::Vtt
        | OutputFormat::Dclx
        | OutputFormat::Chunks => unreachable!("handled above"),
    }
}

fn text_content(value: Option<&str>, name: &str) -> Result<Vec<u8>> {
    value.map(|value| value.as_bytes().to_vec()).ok_or_else(|| {
        PdfConvertError::operation_error(format!("writing {name}"), "output is empty")
    })
}