Skip to main content

cloudiful_docling_convert/
conversion.rs

1use crate::api::{DoclingClient, DoclingConfig};
2use crate::document::{
3    ConvertOptions, GenericFileConvertOptions, InputDocument, InputKind, PdfConvertOptions,
4    TextConvertOptions,
5};
6use crate::error::{PdfConvertError, Result};
7use crate::pdf::PdfInfo;
8use crate::processor::build_pdf_chunk_plan;
9
10#[derive(Debug, Clone)]
11pub struct DoclingRuntimeConfig {
12    pub docling_base_url: String,
13    pub openai_base_url: String,
14    pub vlm_pipeline_model: String,
15    pub picture_description_model: String,
16    pub code_formula_model: String,
17    pub api_key: Option<String>,
18}
19
20impl DoclingRuntimeConfig {
21    pub fn into_docling_config(self) -> DoclingConfig {
22        DoclingConfig {
23            base_url: self.docling_base_url,
24            openai_base_url: self.openai_base_url,
25            vlm_pipeline_model: self.vlm_pipeline_model,
26            picture_description_model: self.picture_description_model,
27            code_formula_model: self.code_formula_model,
28            api_key: self.api_key,
29        }
30    }
31}
32
33#[derive(Debug, Clone)]
34pub struct ConversionBehavior {
35    pub pages_per_file: u32,
36    pub split_input: bool,
37    pub split_by_bookmark: bool,
38    pub chunking: bool,
39    pub batch_size: usize,
40}
41
42impl Default for ConversionBehavior {
43    fn default() -> Self {
44        Self {
45            pages_per_file: 5,
46            split_input: true,
47            split_by_bookmark: false,
48            chunking: false,
49            batch_size: 2,
50        }
51    }
52}
53
54pub fn build_docling_client(config: DoclingRuntimeConfig) -> Result<DoclingClient> {
55    DoclingClient::new(config.into_docling_config())
56}
57
58pub fn build_convert_options(
59    input_kind: InputKind,
60    behavior: &ConversionBehavior,
61) -> Result<ConvertOptions> {
62    match input_kind {
63        InputKind::Pdf => Ok(ConvertOptions::Pdf(PdfConvertOptions {
64            pages_per_file: behavior.pages_per_file,
65            split_input: behavior.split_input,
66            split_by_bookmark: behavior.split_by_bookmark,
67            chunking: behavior.chunking,
68            batch_size: behavior.batch_size,
69        })),
70        InputKind::Docx | InputKind::Markdown => {
71            reject_pdf_only_options(input_kind, behavior)?;
72            Ok(ConvertOptions::Generic(GenericFileConvertOptions {
73                chunking: behavior.chunking,
74            }))
75        }
76        InputKind::Text => {
77            reject_pdf_only_options(input_kind, behavior)?;
78            Ok(ConvertOptions::Text(TextConvertOptions::default()))
79        }
80    }
81}
82
83pub fn build_pdf_options(behavior: &ConversionBehavior) -> Result<PdfConvertOptions> {
84    let options = PdfConvertOptions {
85        pages_per_file: behavior.pages_per_file,
86        split_input: behavior.split_input,
87        split_by_bookmark: behavior.split_by_bookmark,
88        chunking: behavior.chunking,
89        batch_size: behavior.batch_size,
90    };
91    options.validate()?;
92    Ok(options)
93}
94
95pub fn count_input_chunks(input: &InputDocument, behavior: &ConversionBehavior) -> Result<usize> {
96    match input.kind()? {
97        InputKind::Pdf => {
98            let pdf_info = PdfInfo::load_from_bytes(input.bytes.as_ref())?;
99            let total = build_pdf_chunk_plan(&build_pdf_options(behavior)?, &pdf_info)
100                .len()
101                .max(1);
102            Ok(total)
103        }
104        _ => Ok(1),
105    }
106}
107
108fn reject_pdf_only_options(input_kind: InputKind, behavior: &ConversionBehavior) -> Result<()> {
109    if behavior.split_by_bookmark {
110        return Err(PdfConvertError::validation_error(
111            "split_by_bookmark",
112            format!(
113                "bookmark splitting is only available for PDF inputs, got {:?}",
114                input_kind
115            ),
116        ));
117    }
118
119    Ok(())
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn generic_input_rejects_bookmark_splitting() {
128        let err = build_convert_options(
129            InputKind::Docx,
130            &ConversionBehavior {
131                split_by_bookmark: true,
132                ..ConversionBehavior::default()
133            },
134        )
135        .unwrap_err();
136
137        assert!(err.to_string().contains("bookmark splitting"));
138    }
139}