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