Skip to main content

cloudiful_docling_convert/api/
result.rs

1use bytes::{Bytes, BytesMut};
2use futures::{Stream, StreamExt};
3use reqwest::Response;
4use serde_json::Value;
5
6use crate::error::{PdfConvertError, Result};
7use crate::models::{
8    ChunkDocumentResponse, ConvertDocumentResponse, DoclingErrorItem, TaskFailureResult,
9};
10
11use super::transport::handle_response;
12
13#[derive(Debug, Clone)]
14pub enum DoclingResult {
15    Convert(ConvertDocumentResponse),
16    Chunk(ChunkDocumentResponse),
17    Failure(TaskFailureResult),
18    Json(Value),
19    Zip(Bytes),
20}
21
22#[derive(Debug, Clone)]
23pub struct DoclingTaskResult {
24    pub status: crate::models::ConversionStatus,
25    pub result: DoclingResult,
26    pub errors: Vec<String>,
27}
28
29impl DoclingResult {
30    pub fn errors(&self) -> Vec<String> {
31        match self {
32            Self::Convert(response) => response.errors.iter().map(format_error).collect(),
33            Self::Chunk(response) => response
34                .documents
35                .iter()
36                .flat_map(|document| document.errors.iter())
37                .map(format_error)
38                .collect(),
39            Self::Failure(response) => vec![response.failure.message.clone()],
40            Self::Json(value) => collect_json_errors(value),
41            Self::Zip(_) => Vec::new(),
42        }
43    }
44
45    pub fn filename(&self) -> Option<&str> {
46        match self {
47            Self::Convert(response) => Some(&response.document.filename),
48            Self::Chunk(response) => response
49                .chunks
50                .first()
51                .map(|chunk| chunk.filename.as_str())
52                .or_else(|| {
53                    response
54                        .documents
55                        .first()
56                        .map(|document| document.document.filename.as_str())
57                }),
58            Self::Failure(_) | Self::Json(_) | Self::Zip(_) => None,
59        }
60    }
61
62    pub fn into_json(self) -> Option<Value> {
63        match self {
64            Self::Convert(response) => serde_json::to_value(response).ok(),
65            Self::Chunk(response) => serde_json::to_value(response).ok(),
66            Self::Failure(response) => serde_json::to_value(response).ok(),
67            Self::Json(value) => Some(value),
68            Self::Zip(_) => None,
69        }
70    }
71}
72
73pub(crate) async fn parse_response(
74    response: Response,
75    context: &str,
76    body_limit: Option<usize>,
77) -> Result<DoclingResult> {
78    let response = handle_response(response, context).await?;
79    let is_zip = response
80        .headers()
81        .get(reqwest::header::CONTENT_TYPE)
82        .and_then(|value| value.to_str().ok())
83        .is_some_and(|value| value.to_ascii_lowercase().contains("application/zip"));
84    let content_length = response.content_length();
85    let bytes = read_body(response.bytes_stream(), content_length, body_limit, context).await?;
86    parse_bytes(bytes, is_zip)
87}
88
89async fn read_body<S>(
90    mut stream: S,
91    content_length: Option<u64>,
92    body_limit: Option<usize>,
93    context: &str,
94) -> Result<Bytes>
95where
96    S: Stream<Item = std::result::Result<Bytes, reqwest::Error>> + Unpin,
97{
98    if let Some(limit) = body_limit
99        && content_length.is_some_and(|length| length > limit as u64)
100    {
101        return Err(result_too_large(context, limit, content_length));
102    }
103
104    let capacity = body_limit
105        .zip(content_length)
106        .and_then(|(limit, length)| {
107            usize::try_from(length)
108                .ok()
109                .filter(|length| *length <= limit)
110        })
111        .unwrap_or_default();
112    let mut body = BytesMut::with_capacity(capacity);
113    while let Some(chunk) = stream.next().await {
114        let chunk = chunk.map_err(PdfConvertError::from)?;
115        let next_length = body.len().saturating_add(chunk.len());
116        if let Some(limit) = body_limit
117            && next_length > limit
118        {
119            return Err(result_too_large(
120                context,
121                limit,
122                u64::try_from(next_length).ok(),
123            ));
124        }
125        body.extend_from_slice(&chunk);
126    }
127    Ok(body.freeze())
128}
129
130fn result_too_large(context: &str, limit: usize, actual: Option<u64>) -> PdfConvertError {
131    let actual = actual
132        .map(|length| format!("; response is at least {length} bytes"))
133        .unwrap_or_default();
134    PdfConvertError::operation_error(
135        format!("reading {context} response"),
136        format!("result body exceeds maximum of {limit} bytes{actual}"),
137    )
138}
139
140pub(crate) fn parse_bytes(bytes: Bytes, is_zip: bool) -> Result<DoclingResult> {
141    if is_zip || bytes.as_ref().starts_with(b"PK\x03\x04") {
142        return Ok(DoclingResult::Zip(bytes));
143    }
144
145    let value: Value = serde_json::from_slice(&bytes).map_err(|error| {
146        PdfConvertError::parse_error("Docling result response", error.to_string())
147    })?;
148    parse_json(value)
149}
150
151pub(crate) fn parse_json(value: Value) -> Result<DoclingResult> {
152    if value
153        .get("kind")
154        .and_then(Value::as_str)
155        .is_some_and(|kind| kind == "TaskFailureResult")
156        || value.get("failure").is_some()
157    {
158        return serde_json::from_value(value)
159            .map(DoclingResult::Failure)
160            .map_err(|error| {
161                PdfConvertError::parse_error("Docling failure result", error.to_string())
162            });
163    }
164
165    if value.get("chunks").is_some() {
166        return serde_json::from_value(value)
167            .map(DoclingResult::Chunk)
168            .map_err(|error| {
169                PdfConvertError::parse_error("Docling chunk result", error.to_string())
170            });
171    }
172
173    if value.get("document").is_some() && value.get("status").is_some() {
174        return serde_json::from_value(value)
175            .map(DoclingResult::Convert)
176            .map_err(|error| {
177                PdfConvertError::parse_error("Docling conversion result", error.to_string())
178            });
179    }
180
181    Ok(DoclingResult::Json(value))
182}
183
184fn format_error(error: &DoclingErrorItem) -> String {
185    let message = error
186        .error_message
187        .as_deref()
188        .unwrap_or("unknown document error");
189    match error.page_no {
190        Some(page) => format!("page {page}: {message}"),
191        None => message.to_string(),
192    }
193}
194
195fn collect_json_errors(value: &Value) -> Vec<String> {
196    let mut errors = Vec::new();
197    for key in ["error_message", "message"] {
198        if let Some(message) = value.get(key).and_then(Value::as_str) {
199            errors.push(message.to_string());
200        }
201    }
202    if let Some(message) = value
203        .get("failure")
204        .and_then(|failure| failure.get("message"))
205        .and_then(Value::as_str)
206    {
207        errors.push(message.to_string());
208    }
209    if let Some(message) = value.get("detail").and_then(Value::as_str) {
210        errors.push(message.to_string());
211    }
212    if let Some(items) = value.get("errors").and_then(Value::as_array) {
213        errors.extend(items.iter().filter_map(|item| {
214            item.get("error_message")
215                .or_else(|| item.get("message"))
216                .and_then(Value::as_str)
217                .map(ToString::to_string)
218        }));
219    }
220    errors
221}
222
223#[cfg(test)]
224mod tests {
225    use futures::stream;
226
227    use super::*;
228
229    #[tokio::test]
230    async fn limited_reader_accepts_exact_limit() {
231        let chunks = stream::iter([Ok::<_, reqwest::Error>(Bytes::from_static(b"1234"))]);
232
233        let body = read_body(chunks, None, Some(4), "test").await.unwrap();
234
235        assert_eq!(body, Bytes::from_static(b"1234"));
236    }
237
238    #[tokio::test]
239    async fn limited_reader_rejects_chunked_body_over_limit() {
240        let chunks = stream::iter([
241            Ok::<_, reqwest::Error>(Bytes::from_static(b"1234")),
242            Ok::<_, reqwest::Error>(Bytes::from_static(b"5")),
243        ]);
244
245        let error = read_body(chunks, None, Some(4), "test")
246            .await
247            .expect_err("oversized body should fail");
248
249        assert!(error.to_string().contains("exceeds maximum of 4 bytes"));
250    }
251
252    #[tokio::test]
253    async fn limited_reader_rejects_known_length_before_reading() {
254        let chunks = stream::empty::<std::result::Result<Bytes, reqwest::Error>>();
255
256        let error = read_body(chunks, Some(5), Some(4), "test")
257            .await
258            .expect_err("oversized content length should fail");
259
260        assert!(error.to_string().contains("response is at least 5 bytes"));
261    }
262}