diff --git a/litellm-rust/crates/http/tests/websocket.rs b/litellm-rust/crates/http/tests/websocket.rs index 3060556b81b..a5758d45e56 100644 --- a/litellm-rust/crates/http/tests/websocket.rs +++ b/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; diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs index 08df14c988d..e3a4f947433 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs +++ b/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), }, } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index 119a244c6ab..898c24369c3 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/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, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 249c05f4dbc..9a0e0519f12 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/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, }, #[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 { .. } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 4ea6f07e446..7bbc751c614 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/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 diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index f04998434b9..19dd8e24226 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/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), } } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index bd86bab2782..bf0ddc1c7a8 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -302,7 +302,7 @@ async fn prepare_upload_request 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", } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 76601dcd706..31736c7b654 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/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, + #[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::(&error).is_some()); + } + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(error); + assert!(mapped.is_instance_of::(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)]