Skip to main content

cloudiful_docling_convert/api/
docling.rs

1use std::time::{Duration, Instant};
2
3use reqwest::multipart;
4use serde_json::Value;
5
6use crate::document::{ChunkerKind, ChunkingOptions, InputDocument, OutputFormat, PipelineKind};
7use crate::error::{PdfConvertError, Result};
8use crate::models::{ConversionStatus, TaskPostResponse, TaskStatusResponse};
9
10use super::chunk::{build_convert_file_form, build_file_form};
11use super::result::{DoclingResult, DoclingTaskResult, parse_response};
12use super::source::{chunk_source_request, source_request};
13use super::transport::{Transport, default_request_timeout, handle_response, retry_with_backoff};
14
15#[derive(Debug, Clone)]
16pub struct DoclingConfig {
17    pub base_url: String,
18    pub openai_base_url: String,
19    pub vlm_pipeline_model: String,
20    pub picture_description_model: String,
21    pub code_formula_model: String,
22    pub api_key: Option<String>,
23    pub openai_api_key: Option<String>,
24    pub tenant_id: Option<String>,
25    pub request_timeout: Option<Duration>,
26}
27
28impl DoclingConfig {
29    pub fn without_vlm(base_url: impl Into<String>) -> Self {
30        Self {
31            base_url: base_url.into(),
32            openai_base_url: String::new(),
33            vlm_pipeline_model: String::new(),
34            picture_description_model: String::new(),
35            code_formula_model: String::new(),
36            api_key: None,
37            openai_api_key: None,
38            tenant_id: None,
39            request_timeout: None,
40        }
41    }
42}
43
44#[derive(Debug, Clone)]
45pub struct DoclingConvertRequest {
46    pub output_formats: Vec<OutputFormat>,
47    pub page_range: Option<(u32, u32)>,
48    pub chunker: ChunkerKind,
49    pub chunking: ChunkingOptions,
50    pub pipeline: Option<PipelineKind>,
51}
52
53impl DoclingConvertRequest {
54    pub fn for_outputs(output_formats: Vec<OutputFormat>) -> Self {
55        Self {
56            output_formats,
57            page_range: None,
58            chunker: ChunkerKind::None,
59            chunking: ChunkingOptions::hybrid_defaults(),
60            pipeline: None,
61        }
62    }
63
64    pub fn with_chunker(mut self, chunker: ChunkerKind, options: ChunkingOptions) -> Self {
65        self.chunker = chunker;
66        self.chunking = options;
67        self
68    }
69
70    pub fn with_pipeline(mut self, pipeline: Option<PipelineKind>) -> Self {
71        self.pipeline = pipeline;
72        self
73    }
74}
75
76#[derive(Clone)]
77pub struct DoclingClient {
78    transport: Transport,
79}
80
81impl std::fmt::Debug for DoclingClient {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("DoclingClient")
84            .field("base_url", &self.transport.config().base_url)
85            .field("tenant_id", &self.transport.config().tenant_id)
86            .finish_non_exhaustive()
87    }
88}
89
90impl DoclingClient {
91    pub fn new(config: DoclingConfig) -> Result<Self> {
92        Ok(Self {
93            transport: Transport::new(config)?,
94        })
95    }
96
97    pub fn config(&self) -> &DoclingConfig {
98        self.transport.config()
99    }
100
101    pub fn request_timeout(&self) -> Duration {
102        self.config()
103            .request_timeout
104            .unwrap_or_else(default_request_timeout)
105    }
106
107    pub async fn convert_file(
108        &self,
109        input: &InputDocument,
110        request: &DoclingConvertRequest,
111    ) -> Result<DoclingResult> {
112        let operation = || async {
113            let form = self.build_form(input, request)?;
114            let path = match request.chunker {
115                ChunkerKind::None => "convert/file",
116                ChunkerKind::Hybrid => "chunk/hybrid/file",
117                ChunkerKind::Hierarchical => "chunk/hierarchical/file",
118            };
119            let response = self
120                .transport
121                .request(reqwest::Method::POST, path)
122                .multipart(form)
123                .send()
124                .await
125                .map_err(PdfConvertError::from)?;
126            parse_response(response, "Docling file conversion").await
127        };
128
129        retry_with_backoff(operation, "docling_convert_file").await
130    }
131
132    pub async fn submit_file_async(
133        &self,
134        input: &InputDocument,
135        request: &DoclingConvertRequest,
136    ) -> Result<String> {
137        let operation = || async {
138            let form = self.build_form(input, request)?;
139            let path = match request.chunker {
140                ChunkerKind::None => "convert/file/async",
141                ChunkerKind::Hybrid => "chunk/hybrid/file/async",
142                ChunkerKind::Hierarchical => "chunk/hierarchical/file/async",
143            };
144            let response = self
145                .transport
146                .request(reqwest::Method::POST, path)
147                .multipart(form)
148                .send()
149                .await
150                .map_err(PdfConvertError::from)?;
151            let response = handle_response(response, "Docling async submission").await?;
152            let task = response.json::<TaskPostResponse>().await.map_err(|error| {
153                PdfConvertError::parse_error("Docling async submission response", error.to_string())
154            })?;
155            Ok(task.task_id)
156        };
157
158        retry_with_backoff(operation, "docling_submit_file_async").await
159    }
160
161    pub async fn convert_source(
162        &self,
163        url: &str,
164        input_kind: crate::document::InputKind,
165        request: &DoclingConvertRequest,
166    ) -> Result<DoclingResult> {
167        let operation = || async {
168            let path = match request.chunker {
169                ChunkerKind::None => "convert/source",
170                ChunkerKind::Hybrid => "chunk/hybrid/source",
171                ChunkerKind::Hierarchical => "chunk/hierarchical/source",
172            };
173            let body = if request.chunker == ChunkerKind::None {
174                serde_json::to_value(source_request(url, input_kind, request))?
175            } else {
176                serde_json::to_value(chunk_source_request(url, input_kind, request)?)?
177            };
178            let response = self
179                .transport
180                .request(reqwest::Method::POST, path)
181                .json(&body)
182                .send()
183                .await
184                .map_err(PdfConvertError::from)?;
185            parse_response(response, "Docling source conversion").await
186        };
187
188        retry_with_backoff(operation, "docling_convert_source").await
189    }
190
191    pub async fn submit_source_async(
192        &self,
193        url: &str,
194        input_kind: crate::document::InputKind,
195        request: &DoclingConvertRequest,
196    ) -> Result<String> {
197        let operation = || async {
198            let path = match request.chunker {
199                ChunkerKind::None => "convert/source/async",
200                ChunkerKind::Hybrid => "chunk/hybrid/source/async",
201                ChunkerKind::Hierarchical => "chunk/hierarchical/source/async",
202            };
203            let body = if request.chunker == ChunkerKind::None {
204                serde_json::to_value(source_request(url, input_kind, request))?
205            } else {
206                serde_json::to_value(chunk_source_request(url, input_kind, request)?)?
207            };
208            let response = self
209                .transport
210                .request(reqwest::Method::POST, path)
211                .json(&body)
212                .send()
213                .await
214                .map_err(PdfConvertError::from)?;
215            let response = handle_response(response, "Docling source async submission").await?;
216            let task = response.json::<TaskPostResponse>().await.map_err(|error| {
217                PdfConvertError::parse_error(
218                    "Docling source async submission response",
219                    error.to_string(),
220                )
221            })?;
222            Ok(task.task_id)
223        };
224
225        retry_with_backoff(operation, "docling_submit_source_async").await
226    }
227
228    pub async fn wait_for_result(&self, task_id: &str) -> Result<DoclingTaskResult> {
229        self.wait_for_result_with_progress(task_id, |_| async {})
230            .await
231    }
232
233    pub async fn wait_for_result_with_progress<F, Fut>(
234        &self,
235        task_id: &str,
236        mut on_status: F,
237    ) -> Result<DoclingTaskResult>
238    where
239        F: FnMut(TaskStatusResponse) -> Fut + Send,
240        Fut: std::future::Future<Output = ()> + Send,
241    {
242        let deadline = Instant::now() + self.request_timeout();
243        loop {
244            let status = self.poll_task_status(task_id).await?;
245            on_status(status.clone()).await;
246            if status.task_status.is_terminal() {
247                if matches!(
248                    status.task_status,
249                    ConversionStatus::Failure | ConversionStatus::Skipped
250                ) {
251                    return Err(task_failure_error(&status));
252                }
253                let result = self.get_task_result(task_id).await?;
254                return Ok(DoclingTaskResult {
255                    status: status.task_status,
256                    result,
257                    errors: task_status_errors(&status),
258                });
259            }
260
261            if Instant::now() >= deadline {
262                return Err(PdfConvertError::operation_error(
263                    "waiting for Docling task",
264                    format!(
265                        "task {task_id} did not reach a terminal state within {:?}",
266                        self.request_timeout()
267                    ),
268                ));
269            }
270        }
271    }
272
273    pub async fn poll_task_status(&self, task_id: &str) -> Result<TaskStatusResponse> {
274        let operation = || async {
275            let path = format!("status/poll/{task_id}");
276            let response = self
277                .transport
278                .request(reqwest::Method::GET, &format!("{path}?wait=30"))
279                .send()
280                .await
281                .map_err(PdfConvertError::from)?;
282            let response = handle_response(response, "Polling task status").await?;
283            let text = response.text().await.map_err(PdfConvertError::from)?;
284            serde_json::from_str::<TaskStatusResponse>(&text).map_err(|error| {
285                PdfConvertError::parse_error(
286                    "task status response",
287                    format!("task {task_id} returned invalid response: {error}; body: {text}"),
288                )
289            })
290        };
291
292        retry_with_backoff(operation, &format!("check_task_status({task_id})")).await
293    }
294
295    pub async fn check_task_status(&self, task_id: &str) -> Result<bool> {
296        let status = self.poll_task_status(task_id).await?;
297        if matches!(
298            status.task_status,
299            ConversionStatus::Failure | ConversionStatus::Skipped
300        ) {
301            return Err(PdfConvertError::api_task_failed(
302                format!("{:?}", status.task_status),
303                task_status_error(&status),
304            ));
305        }
306        Ok(status.task_status.is_successful())
307    }
308
309    pub async fn get_task_result(&self, task_id: &str) -> Result<DoclingResult> {
310        self.get_task_result_with_connection(task_id, false).await
311    }
312
313    pub async fn get_task_result_with_connection(
314        &self,
315        task_id: &str,
316        close_connection: bool,
317    ) -> Result<DoclingResult> {
318        let operation = || async {
319            let path = format!("result/{task_id}");
320            let mut request = self.transport.request(reqwest::Method::GET, &path);
321            if close_connection {
322                request = request.header(reqwest::header::CONNECTION, "close");
323            }
324            let response = request.send().await.map_err(PdfConvertError::from)?;
325            parse_response(response, "Fetching task result").await
326        };
327
328        retry_with_backoff(operation, &format!("get_task_result({task_id})")).await
329    }
330
331    pub async fn get_task_result_value(&self, task_id: &str) -> Result<Value> {
332        self.get_task_result(task_id)
333            .await?
334            .into_json()
335            .ok_or_else(|| {
336                PdfConvertError::operation_error(
337                    "reading task result",
338                    "task result is a ZIP response and has no JSON value",
339                )
340            })
341    }
342
343    pub(crate) fn build_form(
344        &self,
345        input: &InputDocument,
346        request: &DoclingConvertRequest,
347    ) -> Result<multipart::Form> {
348        match request.chunker {
349            ChunkerKind::None => build_convert_file_form(self, input, request),
350            ChunkerKind::Hybrid | ChunkerKind::Hierarchical => build_file_form(input, request),
351        }
352    }
353}
354
355fn task_status_error(status: &TaskStatusResponse) -> String {
356    status
357        .error_message
358        .clone()
359        .or_else(|| {
360            status
361                .failure
362                .as_ref()
363                .map(|failure| failure.message.clone())
364        })
365        .unwrap_or_else(|| format!("Docling task status is {:?}", status.task_status))
366}
367
368fn task_status_errors(status: &TaskStatusResponse) -> Vec<String> {
369    let mut errors = Vec::new();
370    if let Some(message) = status.error_message.as_deref() {
371        errors.push(message.to_string());
372    }
373    if let Some(failure) = status.failure.as_ref() {
374        errors.push(failure.message.clone());
375    }
376    errors.sort();
377    errors.dedup();
378    errors
379}
380
381fn task_failure_error(status: &TaskStatusResponse) -> PdfConvertError {
382    PdfConvertError::api_task_failed(
383        format!("{:?}", status.task_status),
384        task_status_error(status),
385    )
386}