cloudiful_docling_convert/api/
docling.rs

1use std::time::Duration;
2
3use bytes::Bytes;
4use reqwest::multipart;
5use serde_json::Value;
6
7use crate::document::{InputDocument, OutputFormat};
8use crate::error::{PdfConvertError, Result};
9use crate::models::vlm::{OpenRouterConfigBuilder, VlmConvertOptions};
10use crate::models::{
11    CodeFormulaVlmOptions, PictureDescriptionVlmEngineOptions, TaskPostResponse, TaskStatusResponse,
12};
13
14use super::client::{
15    extract_error_details, get_request, get_request_with_conn_close, handle_response,
16    retry_with_backoff,
17};
18use super::vlm_config::ResolvedVlmConfig;
19
20#[derive(Debug, Clone)]
21pub struct DoclingConfig {
22    pub base_url: String,
23    pub openai_base_url: String,
24    pub vlm_pipeline_model: String,
25    pub picture_description_model: String,
26    pub code_formula_model: String,
27    pub api_key: Option<String>,
28}
29
30#[derive(Debug, Clone)]
31pub struct DoclingConvertRequest {
32    pub output_formats: Vec<OutputFormat>,
33    pub page_range: Option<(u32, u32)>,
34    pub chunking: bool,
35}
36
37impl DoclingConvertRequest {
38    pub fn for_outputs(output_formats: Vec<OutputFormat>) -> Self {
39        Self {
40            output_formats,
41            page_range: None,
42            chunking: false,
43        }
44    }
45}
46
47#[derive(Debug, Clone)]
48pub struct DoclingClient {
49    http_client: reqwest::Client,
50    config: DoclingConfig,
51}
52
53impl DoclingClient {
54    pub fn new(config: DoclingConfig) -> Result<Self> {
55        let http_client = reqwest::Client::builder()
56            .timeout(Duration::from_secs(300))
57            .tcp_keepalive(Duration::from_secs(60))
58            .pool_idle_timeout(Duration::from_secs(30))
59            .build()
60            .map_err(|e| PdfConvertError::api_error(None, e.to_string()))?;
61
62        Ok(Self {
63            http_client,
64            config,
65        })
66    }
67
68    pub fn config(&self) -> &DoclingConfig {
69        &self.config
70    }
71
72    pub async fn convert_file(
73        &self,
74        input: &InputDocument,
75        request: &DoclingConvertRequest,
76    ) -> Result<Value> {
77        let operation = || async {
78            let form = self.build_form(input, request)?;
79            let url = format!("{}/convert/file", self.config.base_url);
80            let response = self
81                .http_client
82                .post(&url)
83                .multipart(form)
84                .send()
85                .await
86                .map_err(PdfConvertError::from)?;
87
88            let response = handle_response(response, "Docling file conversion").await?;
89            response.json::<Value>().await.map_err(|e| {
90                PdfConvertError::parse_error("Docling conversion response", e.to_string())
91            })
92        };
93
94        retry_with_backoff(operation, "docling_convert_file").await
95    }
96
97    pub async fn submit_file_async(
98        &self,
99        input: &InputDocument,
100        request: &DoclingConvertRequest,
101    ) -> Result<String> {
102        let operation = || async {
103            let form = self.build_form(input, request)?;
104            let url = format!("{}/convert/file/async", self.config.base_url);
105            let response = self
106                .http_client
107                .post(&url)
108                .multipart(form)
109                .send()
110                .await
111                .map_err(PdfConvertError::from)?;
112
113            let response = handle_response(response, "Docling async submission").await?;
114            let result = response.json::<TaskPostResponse>().await.map_err(|e| {
115                PdfConvertError::parse_error("Docling async submission response", e.to_string())
116            })?;
117            Ok(result.task_id)
118        };
119
120        retry_with_backoff(operation, "docling_submit_file_async").await
121    }
122
123    pub async fn wait_for_result(&self, task_id: &str) -> Result<Value> {
124        loop {
125            match self.check_task_status(task_id).await? {
126                true => {
127                    tokio::time::sleep(Duration::from_millis(500)).await;
128                    return self.get_task_result(task_id, true).await;
129                }
130                false => tokio::time::sleep(Duration::from_secs(5)).await,
131            }
132        }
133    }
134
135    pub async fn check_task_status(&self, task_id: &str) -> Result<bool> {
136        let operation = || async {
137            let url = format!("{}/status/poll/{}", self.config.base_url, task_id);
138            let response = get_request(&self.http_client, &url, "Polling task status").await?;
139            let response_text = response.text().await?;
140
141            match serde_json::from_str::<TaskStatusResponse>(&response_text) {
142                Ok(status_response) => match status_response.task_status.as_str() {
143                    "success" => Ok(true),
144                    "failure" | "revoked" => {
145                        let error_details = extract_error_details(&response_text);
146                        Err(PdfConvertError::api_task_failed(
147                            status_response.task_status,
148                            error_details,
149                        ))
150                    }
151                    _ => Ok(false),
152                },
153                Err(_) => Err(PdfConvertError::parse_error(
154                    "task status response",
155                    format!(
156                        "Task {} returned invalid response: {}",
157                        task_id, response_text
158                    ),
159                )),
160            }
161        };
162
163        retry_with_backoff(operation, &format!("check_task_status({task_id})")).await
164    }
165
166    pub async fn get_task_result(&self, task_id: &str, use_new_conn: bool) -> Result<Value> {
167        let operation = || async {
168            let url = format!("{}/result/{}", self.config.base_url, task_id);
169            let response = get_request_with_conn_close(
170                &self.http_client,
171                &url,
172                "Fetching task result",
173                use_new_conn,
174            )
175            .await?;
176
177            response
178                .json::<Value>()
179                .await
180                .map_err(|e| PdfConvertError::parse_error("task result response", e.to_string()))
181        };
182
183        retry_with_backoff(operation, &format!("get_task_result({task_id})")).await
184    }
185
186    fn build_form(
187        &self,
188        input: &InputDocument,
189        request: &DoclingConvertRequest,
190    ) -> Result<multipart::Form> {
191        let input_kind = input.kind()?;
192        let part = multipart::Part::stream(reqwest::Body::from(Bytes::clone(&input.bytes)))
193            .file_name(input.filename.clone())
194            .mime_str(&input.media_type)
195            .map_err(|e| {
196                PdfConvertError::api_error(None, format!("Failed to create multipart part: {e}"))
197            })?;
198
199        let mut form = multipart::Form::new()
200            .part("files", part)
201            .text("from_formats", input_kind.from_formats_value().to_string())
202            .text("target_type", "inbody");
203
204        for format in &request.output_formats {
205            form = form.text("to_formats", format.as_api_value().to_string());
206        }
207
208        if let Some((start_page, end_page)) = request.page_range {
209            form = form.text("page_range", start_page.to_string());
210            form = form.text("page_range", end_page.to_string());
211        }
212
213        if request.chunking {
214            form = form.text("include_chunking", "true");
215        }
216
217        if input_kind.supports_vlm()
218            && let Some(vlm_config) = self.config.resolved_vlm_config()?
219        {
220            form = self.apply_vlm_config(form, &vlm_config)?;
221        }
222
223        Ok(form)
224    }
225
226    fn apply_vlm_config(
227        &self,
228        mut form: multipart::Form,
229        vlm_config: &ResolvedVlmConfig,
230    ) -> Result<multipart::Form> {
231        let picture_description_custom_config =
232            PictureDescriptionVlmEngineOptions::for_openai_compatible(
233                &vlm_config.openai_base_url,
234                &vlm_config.api_key,
235                &vlm_config.picture_description_model,
236                "Describe this image in a few sentences.",
237                300,
238                60,
239            );
240        let code_formula_custom_config = CodeFormulaVlmOptions {
241            scale: Some(2.0),
242            max_size: None,
243            extract_code: Some(true),
244            extract_formulas: Some(true),
245            engine_options: OpenRouterConfigBuilder::engine_options(
246                &vlm_config.openai_base_url,
247                &vlm_config.api_key,
248                &vlm_config.code_formula_model,
249                30,
250                2,
251            ),
252            model_spec: OpenRouterConfigBuilder::model_spec(
253                &vlm_config.code_formula_model,
254                "Recognize code blocks and mathematical formulas in the image. For code, output the full code; for mathematical formulas, output in LaTeX format.",
255                1000,
256            ),
257        };
258        let vlm_pipeline_custom_config = VlmConvertOptions {
259            engine_options: OpenRouterConfigBuilder::engine_options(
260                &vlm_config.openai_base_url,
261                &vlm_config.api_key,
262                &vlm_config.vlm_pipeline_model,
263                30,
264                2,
265            ),
266            model_spec: OpenRouterConfigBuilder::model_spec(
267                &vlm_config.vlm_pipeline_model,
268                "",
269                1000,
270            ),
271            scale: Some(1.0),
272            max_size: None,
273            batch_size: None,
274            force_backend_text: true,
275        };
276
277        form = form.text(
278            "vlm_pipeline_custom_config",
279            serde_json::to_string(&vlm_pipeline_custom_config)?,
280        );
281        form = form.text(
282            "picture_description_custom_config",
283            serde_json::to_string(&picture_description_custom_config)?,
284        );
285        form = form.text(
286            "code_formula_custom_config",
287            serde_json::to_string(&code_formula_custom_config)?,
288        );
289        form = form.text("do_code_enrichment", "True");
290        form = form.text("do_formula_enrichment", "True");
291        form = form.text("do_picture_description", "True");
292        form = form.text("ocr_engine", "rapidocr");
293        form = form.text("image_export_mode", "placeholder");
294
295        Ok(form)
296    }
297}
298
299#[cfg(test)]
300#[path = "docling_tests.rs"]
301mod tests;
cloudiful_docling_convert/api/docling.rs

cloudiful_docling_convert/api/
docling.rs
