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