Skip to main content

cloudiful_docling_convert/processor/service/
output.rs

1use std::path::{Path, PathBuf};
2
3use serde_json::Value;
4use tokio::fs;
5
6use super::DocumentConverter;
7use crate::document::{
8    ConvertedChunk, ConvertedDocument, ConvertedDocumentMetadata, InputDocument, InputKind,
9    OutputFormat,
10};
11use crate::error::{PdfConvertError, Result};
12use crate::models::{Bookmark, ChunkMetadata};
13
14impl DocumentConverter {
15    pub fn calculate_output_path(
16        output_dir: &Path,
17        filename: &str,
18        output_format: OutputFormat,
19    ) -> PathBuf {
20        let stem = Path::new(filename)
21            .file_stem()
22            .and_then(|value| value.to_str())
23            .unwrap_or("output");
24        output_dir.join(format!("{}.{}", stem, output_format.extension()))
25    }
26
27    pub(super) async fn write_output_file(
28        output_path: &Path,
29        document: &ConvertedDocument,
30        output_format: OutputFormat,
31    ) -> Result<()> {
32        if let Some(parent) = output_path.parent() {
33            fs::create_dir_all(parent).await.map_err(|e| {
34                PdfConvertError::io_error(
35                    format!("creating output directory: {}", parent.display()),
36                    e,
37                )
38            })?;
39        }
40
41        let content = Self::select_output_content(document, output_format)?;
42        fs::write(output_path, content).await.map_err(|e| {
43            PdfConvertError::io_error(format!("writing output file: {}", output_path.display()), e)
44        })?;
45
46        Ok(())
47    }
48
49    fn select_output_content(
50        document: &ConvertedDocument,
51        output_format: OutputFormat,
52    ) -> Result<Vec<u8>> {
53        match output_format {
54            OutputFormat::Md => document
55                .markdown
56                .as_ref()
57                .map(|value| value.as_bytes().to_vec())
58                .ok_or_else(|| {
59                    PdfConvertError::operation_error("writing markdown", "markdown output is empty")
60                }),
61            OutputFormat::Text => document
62                .text
63                .as_ref()
64                .map(|value| value.as_bytes().to_vec())
65                .ok_or_else(|| {
66                    PdfConvertError::operation_error("writing text", "text output is empty")
67                }),
68            OutputFormat::Html => document
69                .html
70                .as_ref()
71                .map(|value| value.as_bytes().to_vec())
72                .ok_or_else(|| {
73                    PdfConvertError::operation_error("writing html", "html output is empty")
74                }),
75            OutputFormat::Doctags => document
76                .doctags
77                .as_ref()
78                .map(|value| value.as_bytes().to_vec())
79                .ok_or_else(|| {
80                    PdfConvertError::operation_error("writing doctags", "doctags output is empty")
81                }),
82            OutputFormat::Json => document
83                .json
84                .as_ref()
85                .map(serde_json::to_vec_pretty)
86                .transpose()
87                .map_err(PdfConvertError::from)?
88                .ok_or_else(|| {
89                    PdfConvertError::operation_error("writing json", "json output is empty")
90                }),
91        }
92    }
93
94    pub(super) fn chunk_from_result(
95        metadata: Option<ChunkMetadata>,
96        raw_result: Value,
97    ) -> ConvertedChunk {
98        ConvertedChunk {
99            metadata,
100            markdown: extract_document_field(&raw_result, "md_content"),
101            text: extract_document_field(&raw_result, "text_content"),
102            json: raw_result
103                .get("document")
104                .and_then(|document| document.get("json_content"))
105                .cloned(),
106            html: extract_document_field(&raw_result, "html_content"),
107            doctags: extract_document_field(&raw_result, "doctags_content"),
108            raw_result,
109        }
110    }
111
112    pub(super) fn assemble_document(
113        input: &InputDocument,
114        input_kind: InputKind,
115        page_count: Option<u32>,
116        outlines: Vec<Bookmark>,
117        chunks: Vec<ConvertedChunk>,
118    ) -> ConvertedDocument {
119        let markdown = join_optional_chunks(chunks.iter().map(|chunk| chunk.markdown.as_deref()));
120        let text = join_optional_chunks(chunks.iter().map(|chunk| chunk.text.as_deref()));
121        let html = join_optional_chunks(chunks.iter().map(|chunk| chunk.html.as_deref()));
122        let doctags = join_optional_chunks(chunks.iter().map(|chunk| chunk.doctags.as_deref()));
123        let json_values: Vec<Value> = chunks
124            .iter()
125            .filter_map(|chunk| chunk.json.clone())
126            .collect();
127        let json = match json_values.len() {
128            0 => None,
129            1 => json_values.into_iter().next(),
130            _ => Some(Value::Array(json_values)),
131        };
132
133        ConvertedDocument {
134            filename: input.filename.clone(),
135            markdown,
136            text,
137            json,
138            html,
139            doctags,
140            chunks,
141            metadata: ConvertedDocumentMetadata {
142                input_kind,
143                media_type: input.media_type.clone(),
144                page_count,
145                outlines,
146            },
147            errors: Vec::new(),
148        }
149    }
150}
151
152fn extract_document_field(raw_result: &Value, field: &str) -> Option<String> {
153    raw_result
154        .get("document")
155        .and_then(|document| document.get(field))
156        .and_then(|value| value.as_str())
157        .map(ToString::to_string)
158}
159
160fn join_optional_chunks<'a>(values: impl Iterator<Item = Option<&'a str>>) -> Option<String> {
161    let mut collected = Vec::new();
162    for value in values.flatten() {
163        collected.push(value.trim_end_matches('\n').to_string());
164    }
165
166    if collected.is_empty() {
167        None
168    } else {
169        Some(collected.join("\n"))
170    }
171}