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 02d2ccbdeea..9bd29ce601f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,7 +1,6 @@ use litellm_core::ocr::Error; use pyo3::exceptions::{PyFileNotFoundError, PyOSError}; use pyo3::prelude::*; -use pyo3::types::PyDict; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; @@ -42,15 +41,8 @@ fn upstream_error( body: String, headers: Vec<(String, String)>, ) -> PyResult { - let kwargs = PyDict::new(py); - kwargs.set_item("content", &body)?; - kwargs.set_item("headers", headers)?; - let response = py - .import("httpx")? - .getattr("Response")? - .call((status,), Some(&kwargs))?; let error = RustUpstreamError::new_err((status, body)); - error.value(py).setattr("response", response)?; + error.value(py).setattr("headers", headers)?; Ok(error) } @@ -89,9 +81,15 @@ mod tests { let mapped = to_pyerr(Error::Provider { status: 429, body: r#"{"message":"rate limited"}"#.to_string(), - headers: Vec::new(), + headers: vec![("Retry-After".to_string(), "17".to_string())], }); assert!(mapped.is_instance_of::(py)); + let headers: Vec<(String, String)> = mapped + .value(py) + .getattr("headers") + .and_then(|headers| headers.extract()) + .expect("OCR failures retain provider headers"); + assert_eq!(headers, vec![("Retry-After".to_string(), "17".to_string())]); let args: (u16, String) = mapped .value(py) .getattr("args") diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr_lifecycle.py index 1958fdf8cf3..e22722d22c4 100644 --- a/litellm/rust_bridge/ocr_lifecycle.py +++ b/litellm/rust_bridge/ocr_lifecycle.py @@ -4,6 +4,7 @@ from collections.abc import Awaitable, Mapping, Sequence from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -41,6 +42,27 @@ def _binding(value: object) -> NativeOcrLifecycle | None: NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding) +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) + + +class UpstreamFailure(Exception): + def __init__(self, response: httpx.Response, cause: Exception) -> None: + super().__init__(str(cause)) + self.message: Final = str(cause) + self.response: Final = response + self.status_code: Final = response.status_code + self.__cause__ = cause + + +def _upstream_failure(error: Exception) -> Exception: + try: + status, body = _UPSTREAM_ARGS.validate_python(error.args) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return error + return UpstreamFailure(httpx.Response(status, content=body.encode(), headers=headers), error) + def select(request: LiteLLMOcrRequest) -> NativeOcrLifecycle | None: if request.kwargs.get("aocr"): @@ -63,18 +85,18 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper ExceptionMapper, litellm.exception_type ) + original: Final = _upstream_failure(error) try: return mapper( model=model, custom_llm_provider=request_provider, - original_exception=error, + original_exception=original, completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs ) except Exception as public_error: - response: Final = getattr(error, "response", None) - if isinstance(response, httpx.Response): - public_error.response = response - public_error.status_code = response.status_code + if isinstance(original, UpstreamFailure): + public_error.response = original.response + public_error.status_code = original.status_code public_error.__context__ = error return public_error diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 3e95258fb36..e360401a435 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -13,6 +13,7 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, Re from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, + call_native, call_native_aocr, call_native_ocr, ) @@ -43,10 +44,7 @@ async def test_ocr_contract_upstream_status( "num_retries": 0, } with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, **arguments) - else: - call_native_ocr(ocr_server, **arguments) + await call_native(ocr_server, asynchronous, **arguments) assert caught.value.status_code == upstream.status assert caught.value.response.status_code == upstream.status @@ -64,10 +62,7 @@ async def test_ocr_contract_provider_error_details( headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"} ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers)) with pytest.raises(litellm.RateLimitError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, num_retries=0) - else: - call_native_ocr(ocr_server, num_retries=0) + await call_native(ocr_server, asynchronous, num_retries=0) response: Final = caught.value.response assert isinstance(response, httpx.Response) if preserved == "body": @@ -87,10 +82,7 @@ async def test_ocr_contract_invalid_response_format( ) -> None: ocr_server.expected_requests = 0 with pytest.raises(litellm.UnsupportedParamsError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, req_format="bogus", num_retries=0) - else: - call_native_ocr(ocr_server, req_format="bogus", num_retries=0) + await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0) assert caught.value.status_code == 400 for value in ("req_format", "bogus", "native", "litellm"): assert value in str(caught.value) @@ -116,10 +108,7 @@ async def test_ocr_contract_malformed_document_is_actionable( ) -> None: ocr_server.expected_requests = None with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, document=document, num_retries=0) - else: - call_native_ocr(ocr_server, document=document, num_retries=0) + await call_native(ocr_server, asynchronous, document=document, num_retries=0) assert caught.value.status_code == 400 assert field.lower() in str(caught.value).lower() assert "NoneType: None" not in str(caught.value) @@ -141,10 +130,7 @@ async def test_ocr_contract_azure_invalid_options_are_bad_requests( ocr_server.expected_requests = 0 arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0} with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, **arguments) - else: - call_native_ocr(ocr_server, **arguments) + await call_native(ocr_server, asynchronous, **arguments) assert caught.value.status_code == 400 assert field in str(caught.value) assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index 7114e42a59e..b60cf5eac02 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -42,6 +42,10 @@ async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResp return await call_aocr(server, **kwargs) +async def call_native(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse: + return await call_native_aocr(server, **kwargs) if asynchronous else call_native_ocr(server, **kwargs) + + def request_body(kwargs: dict[str, object]) -> dict[str, object]: additional_args = kwargs["additional_args"] assert isinstance(additional_args, dict)