fix(rust): preserve OCR download status and reduce error size

litellm-rust/crates/http/tests/websocket.rs

@@ -289,7 +289,7 @@ async fn websocket_uses_system_roots(
         let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await;
         match result {
             Ok(mut socket) if connects => round_trip(&mut socket).await,
-            Err(Error::Connect(_)) if !connects => {}
+            Err(error) if !connects && error.is_connect() => {}
             other => panic!("unexpected handshake outcome: {:?}", other.map(drop)),
         }
         return;

litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs

@@ -332,7 +332,7 @@ pub(super) fn error_class(body: String, status: u16, headers: reqwest::header::H
         Ok(_) | Err(_) => Error::Provider {
             status,
             body,
-            headers,
+            headers: Box::new(headers),
         },
     }
 }

litellm-rust/crates/llms/src/base_llm/ocr/document.rs

@@ -68,7 +68,8 @@ pub async fn inline_remote_document(
                 max_redirects: OCR_MAX_FETCH_REDIRECTS,
             },
         )
-        .await?;
+        .await
+        .map_err(Error::DocumentDownload)?;
     let result = document.with_source(format!(
         "data:{};base64,{}",
         downloaded.content_type,

litellm-rust/crates/llms/src/base_llm/ocr/error.rs

@@ -12,7 +12,7 @@ pub enum Error {
     Provider {
         status: u16,
         body: String,
-        headers: HeaderMap,
+        headers: Box<HeaderMap>,
     },
     #[error("File is empty or could not be read")]
     EmptyFile,
@@ -116,6 +116,8 @@ pub enum Error {
     Headers(#[from] litellm_http::request::HeaderError),
     #[error(transparent)]
     Http(#[from] litellm_http::Error),
+    #[error(transparent)]
+    DocumentDownload(litellm_http::Error),
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
@@ -174,6 +176,14 @@ impl Error {
                 body,
                 headers,
             } => Some((*status, body, headers)),
+            Self::DocumentDownload(error) => error.status().map(|status| {
+                let body = if error.is_timeout() {
+                    "OCR document download timed out"
+                } else {
+                    "OCR document download failed"
+                };
+                (status, body, &*NO_HEADERS)
+            }),
             Self::Http(error) => error
                 .upstream_response()
                 .map(|(status, body)| (status, body, &*NO_HEADERS)),
@@ -253,7 +263,8 @@ impl Error {
             | Self::HostCall
             | Self::Params(_)
             | Self::Headers(_)
-            | Self::Http(_) => false,
+            | Self::Http(_)
+            | Self::DocumentDownload(_) => false,
         }
     }
 
@@ -295,7 +306,8 @@ impl Error {
             | Self::Auth(_)
             | Self::Params(_)
             | Self::Headers(_)
-            | Self::Http(_) => false,
+            | Self::Http(_)
+            | Self::DocumentDownload(_) => false,
         }
     }
 
@@ -332,7 +344,7 @@ impl Error {
             | Self::Params(_)
             | Self::Headers(_) => true,
             Self::Auth(error) => error.is_request(),
-            Self::Http(error) => error.is_request(),
+            Self::Http(error) | Self::DocumentDownload(error) => error.is_request(),
             Self::Provider { .. }
             | Self::FileRead { .. }
             | Self::DocumentTask(_)
@@ -364,7 +376,7 @@ impl Error {
             | Self::InvalidPollUrl { .. }
             | Self::PollOrigin
             | Self::PollTimeout => true,
-            Self::Http(error) => error.is_response(),
+            Self::Http(error) | Self::DocumentDownload(error) => error.is_response(),
             Self::Provider { .. }
             | Self::EmptyFile
             | Self::FileRead { .. }

litellm-rust/crates/llms/src/base_llm/ocr/handler.rs

@@ -173,7 +173,7 @@ pub async fn read_response_bytes(
         return Err(Error::Provider {
             status,
             body,
-            headers,
+            headers: Box::new(headers),
         });
     }
     if response

litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs

@@ -366,7 +366,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static {
         Error::Provider {
             status: status_code,
             body: error_message,
-            headers,
+            headers: Box::new(headers),
         }
     }
 

litellm-rust/crates/llms/src/reducto/ocr/transformation.rs

@@ -302,7 +302,7 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
 fn unsupported_source() -> Error {
     Error::UnsupportedDocument {
         provider: "Reducto",
-        guidance: "pass a reducto:// id or a data URI, plain HTTP URLs must be uploaded first",
+        guidance: "pass a reducto:// id or a data URI, or upload the file first for plain HTTP URLs",
     }
 }
 

litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs

@@ -80,10 +80,10 @@ mod tests {
         Error::Provider {
             status: 429,
             body: r#"{"message":"rate limited"}"#.into(),
-            headers: reqwest::header::HeaderMap::from_iter([(
+            headers: Box::new(reqwest::header::HeaderMap::from_iter([(
                 reqwest::header::RETRY_AFTER,
                 reqwest::header::HeaderValue::from_static("17"),
-            )]),
+            )])),
         },
         vec![("retry-after".into(), "17".into())]
     )]
@@ -105,6 +105,68 @@ mod tests {
         });
     }
 
+    #[rstest]
+    #[case::not_found(Some(404), 404, "OCR document download failed")]
+    #[case::rate_limited(Some(429), 429, "OCR document download failed")]
+    #[case::timeout(None, 408, "OCR document download timed out")]
+    #[tokio::test]
+    async fn document_download_failures_keep_the_upstream_exception(
+        #[case] response_status: Option<u16>,
+        #[case] expected_status: u16,
+        #[case] expected_body: &str,
+    ) {
+        use std::time::Duration;
+
+        use litellm_http::{Client, media::MediaFetcher};
+        use litellm_llms::base_llm::ocr::{
+            document::inline_remote_document, transformation::OcrConnection,
+        };
+        use litellm_llms_types::formats::ocr::OcrDocument;
+        use tokio::{
+            io::{AsyncReadExt, AsyncWriteExt},
+            net::TcpListener,
+        };
+
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(async move {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut request = [0_u8; 1];
+            socket.read_exact(&mut request).await.unwrap();
+            if let Some(status) = response_status {
+                socket.write_all(format!("HTTP/1.1 {status} Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").as_bytes()).await.unwrap();
+            } else {
+                std::future::pending::<()>().await;
+            }
+        });
+        let connection = OcrConnection {
+            timeout: Duration::from_millis(100),
+            ..OcrConnection::default()
+        };
+        let error = inline_remote_document(
+            &MediaFetcher::for_test(Client::for_test()),
+            OcrDocument::DocumentUrl {
+                document_url: format!("http://{address}/document.pdf"),
+                extra_fields: Default::default(),
+            },
+            &connection,
+        )
+        .await
+        .unwrap_err();
+        server.abort();
+        if response_status.is_none() {
+            assert!(litellm_error::find::<tokio::time::error::Elapsed>(&error).is_some());
+        }
+        Python::initialize();
+        Python::attach(|py| {
+            let mapped = to_pyerr(error);
+            assert!(mapped.is_instance_of::<RustUpstreamError>(py), "{mapped}");
+            let args: (u16, String) = mapped.value(py).getattr("args").unwrap().extract().unwrap();
+            assert_eq!(args, (expected_status, expected_body.into()));
+            assert_eq!(status_code(py, &mapped), Some(expected_status));
+        });
+    }
+
     #[rstest]
     #[case::missing_document_url(Error::MissingDocumentUrl)]
     #[case::invalid_features(Error::Features)]