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
203 for format in &request.output_formats {
204 form = form.text("to_formats", format.as_api_value().to_string());
205 }
206
207 if let Some((start_page, end_page)) = request.page_range {
208 form = form.text("page_range", start_page.to_string());
209 form = form.text("page_range", end_page.to_string());
210 }
211
212 if request.chunking {
213 form = form.text("include_chunking", "true");
214 }
215
216 if input_kind.supports_vlm() {
217 if let Some(vlm_config) = self.config.resolved_vlm_config()? {
218 form = self.apply_vlm_config(form, &vlm_config)?;
219 }
220 }
221
222 Ok(form)
223 }
224
225 fn apply_vlm_config(
226 &self,
227 mut form: multipart::Form,
228 vlm_config: &ResolvedVlmConfig,
229 ) -> Result<multipart::Form> {
230 let picture_description_custom_config =
231 PictureDescriptionVlmEngineOptions::for_openai_compatible(
232 &vlm_config.openai_base_url,
233 &vlm_config.api_key,
234 &vlm_config.picture_description_model,
235 "Describe this image in a few sentences.",
236 300,
237 60,
238 );
239 let code_formula_custom_config = CodeFormulaVlmOptions {
240 scale: Some(2.0),
241 max_size: None,
242 extract_code: Some(true),
243 extract_formulas: Some(true),
244 engine_options: OpenRouterConfigBuilder::engine_options(
245 &vlm_config.openai_base_url,
246 &vlm_config.api_key,
247 &vlm_config.code_formula_model,
248 30,
249 2,
250 ),
251 model_spec: OpenRouterConfigBuilder::model_spec(
252 &vlm_config.code_formula_model,
253 "Recognize code blocks and mathematical formulas in the image. For code, output the full code; for mathematical formulas, output in LaTeX format.",
254 1000,
255 ),
256 };
257 let vlm_pipeline_custom_config = VlmConvertOptions {
258 engine_options: OpenRouterConfigBuilder::engine_options(
259 &vlm_config.openai_base_url,
260 &vlm_config.api_key,
261 &vlm_config.vlm_pipeline_model,
262 30,
263 2,
264 ),
265 model_spec: OpenRouterConfigBuilder::model_spec(
266 &vlm_config.vlm_pipeline_model,
267 "",
268 1000,
269 ),
270 scale: Some(1.0),
271 max_size: None,
272 batch_size: None,
273 force_backend_text: true,
274 };
275
276 form = form.text(
277 "vlm_pipeline_custom_config",
278 serde_json::to_string(&vlm_pipeline_custom_config)?,
279 );
280 form = form.text(
281 "picture_description_custom_config",
282 serde_json::to_string(&picture_description_custom_config)?,
283 );
284 form = form.text(
285 "code_formula_custom_config",
286 serde_json::to_string(&code_formula_custom_config)?,
287 );
288 form = form.text("do_code_enrichment", "True");
289 form = form.text("do_formula_enrichment", "True");
290 form = form.text("do_picture_description", "True");
291 form = form.text("ocr_engine", "rapidocr");
292 form = form.text("image_export_mode", "placeholder");
293
294 Ok(form)
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use super::*;
301 use crate::document::InputKind;
302
303 #[test]
304 fn build_form_uses_input_media_type_and_format() {
305 let client = DoclingClient::new(DoclingConfig {
306 base_url: "http://localhost:5001/v1".to_string(),
307 openai_base_url: "http://localhost:1234/v1".to_string(),
308 vlm_pipeline_model: "vlm".to_string(),
309 picture_description_model: "pic".to_string(),
310 code_formula_model: "code".to_string(),
311 api_key: Some("secret".to_string()),
312 })
313 .unwrap();
314
315 let input = InputDocument::new("notes.md", "text/markdown", Bytes::from_static(b"# hello"));
316 let request = DoclingConvertRequest {
317 output_formats: vec![OutputFormat::Md, OutputFormat::Text],
318 page_range: None,
319 chunking: false,
320 };
321
322 let form = client.build_form(&input, &request).unwrap();
323 let debug = format!("{form:?}");
324 assert!(debug.contains("text/markdown"));
325 assert!(debug.contains("notes.md"));
326 assert!(debug.contains("to_formats"));
327 }
328
329 #[test]
330 fn build_form_skips_page_range_for_generic_requests() {
331 let client = DoclingClient::new(DoclingConfig {
332 base_url: "http://localhost:5001/v1".to_string(),
333 openai_base_url: "http://localhost:1234/v1".to_string(),
334 vlm_pipeline_model: "vlm".to_string(),
335 picture_description_model: "pic".to_string(),
336 code_formula_model: "code".to_string(),
337 api_key: Some("secret".to_string()),
338 })
339 .unwrap();
340
341 let input = InputDocument::new(
342 "doc.docx",
343 "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
344 Bytes::from_static(b"PK"),
345 );
346 let request = DoclingConvertRequest::for_outputs(vec![OutputFormat::Md]);
347
348 let form = client.build_form(&input, &request).unwrap();
349 let debug = format!("{form:?}");
350 assert!(!debug.contains("page_range"));
351 assert!(debug.contains("from_formats"));
352 }
353
354 #[test]
355 fn build_form_uses_generic_first_wave_formats() {
356 let client =
357 DoclingClient::new(DoclingConfig::without_vlm("http://localhost:5001/v1")).unwrap();
358
359 let input = InputDocument::new("sheet.csv", "text/csv", Bytes::from_static(b"a,b"));
360 let request = DoclingConvertRequest::for_outputs(vec![OutputFormat::Html]);
361
362 let form = client.build_form(&input, &request).unwrap();
363 let debug = format!("{form:?}");
364 assert!(debug.contains("text/csv"));
365 assert!(debug.contains("sheet.csv"));
366 assert!(debug.contains("csv"));
367 assert!(debug.contains("to_formats"));
368 }
369
370 #[test]
371 fn build_form_supports_override_only_second_wave_formats() {
372 let client =
373 DoclingClient::new(DoclingConfig::without_vlm("http://localhost:5001/v1")).unwrap();
374
375 let input = InputDocument::new(
376 "filing.xml",
377 "application/xml",
378 Bytes::from_static(b"<article />"),
379 )
380 .with_input_kind(InputKind::XmlJats);
381 let request = DoclingConvertRequest::for_outputs(vec![OutputFormat::Md]);
382
383 let form = client.build_form(&input, &request).unwrap();
384 let debug = format!("{form:?}");
385 assert_eq!(input.kind().unwrap(), InputKind::XmlJats);
386 assert!(debug.contains("application/xml"));
387 assert!(debug.contains("from_formats"));
388 }
389
390 #[test]
391 fn build_form_skips_vlm_fields_when_runtime_config_is_missing() {
392 let client =
393 DoclingClient::new(DoclingConfig::without_vlm("http://localhost:5001/v1")).unwrap();
394
395 let input = InputDocument::new("notes.md", "text/markdown", Bytes::from_static(b"# hello"));
396 let request = DoclingConvertRequest::for_outputs(vec![OutputFormat::Md]);
397
398 let form = client.build_form(&input, &request).unwrap();
399 let debug = format!("{form:?}");
400 assert!(!debug.contains("vlm_pipeline_custom_config"));
401 assert!(!debug.contains("picture_description_custom_config"));
402 assert!(!debug.contains("code_formula_custom_config"));
403 assert!(!debug.contains("do_code_enrichment"));
404 assert!(!debug.contains("do_formula_enrichment"));
405 assert!(!debug.contains("do_picture_description"));
406 }
407}