make test pass

This commit is contained in:
Yujong Lee 2026-09-17 08:55:59 -07:00
parent 370cdaabf9
commit f1ea94fee7
11 changed files with 182 additions and 126 deletions

View file

@ -133,14 +133,9 @@ pub async fn read_json_response<T: DeserializeOwned>(
pub(crate) async fn read_response_bytes(
mut response: reqwest::Response,
max_response_bytes: usize,
limit: usize,
) -> Result<Bytes, crate::ocr::Error> {
let status = response.status();
let limit = if status.is_success() {
max_response_bytes
} else {
max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1))
};
if status.is_success()
&& response
.content_length()
@ -162,7 +157,7 @@ pub(crate) async fn read_response_bytes(
if !status.is_success() {
return Err(crate::transport::Error::Http {
status: status.as_u16(),
body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)),
body: String::from_utf8_lossy(&bytes).into_owned(),
}
.into());
}

View file

@ -429,7 +429,7 @@ mod tests {
client.document_fetcher(),
OcrDocument::ImageUrl {
image_url: format!("http://{address}/image"),
extra_fields: Map::from_iter([("detail".into(), "high".into())]),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
},
&OcrConnection::default(),
)
@ -441,7 +441,7 @@ mod tests {
converted,
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::from_iter([("detail".into(), "high".into())]),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
}
);
assert!(!request.to_ascii_lowercase().contains("authorization"));

View file

@ -113,7 +113,6 @@ impl From<crate::call_arguments::ArgumentError> for Error {
impl Error {
pub fn http_status_code(&self) -> Option<u16> {
match self {
Self::MissingDocumentUrl => Some(500),
Self::Provider { status, .. }
| Self::Transport(crate::transport::Error::Http { status, .. }) => Some(*status),
error if error.is_request() => Some(400),
@ -142,6 +141,7 @@ impl Error {
| Self::Features
| Self::DotModel
| Self::InvalidRequest(_)
| Self::InvalidProvider(_)
| Self::Params(_)
| Self::Headers(_)
)

View file

@ -22,13 +22,13 @@ pub enum OcrDocument {
DocumentUrl {
document_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, String>,
extra_fields: BTreeMap<String, Option<String>>,
},
#[serde(rename = "image_url")]
ImageUrl {
image_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, String>,
extra_fields: BTreeMap<String, Option<String>>,
},
}
@ -720,45 +720,24 @@ mod tests {
}
}
#[test]
fn document_variants_preserve_provider_fields_when_rewriting_sources() {
for (value, original, replacement, expected) in [
(
json!({
"type":"document_url",
"document_url":"https://example.com/input.pdf",
"document_name":"input.pdf"
}),
"https://example.com/input.pdf",
"data:application/pdf;base64,AA==",
json!({
"type":"document_url",
"document_url":"data:application/pdf;base64,AA==",
"document_name":"input.pdf"
}),
),
(
json!({
"type":"image_url",
"image_url":"https://example.com/input.png",
"detail":"high"
}),
"https://example.com/input.png",
"data:image/png;base64,AA==",
json!({
"type":"image_url",
"image_url":"data:image/png;base64,AA==",
"detail":"high"
}),
),
] {
let document: OcrDocument = serde_json::from_value(value).unwrap();
assert_eq!(document.source(), original);
assert_eq!(
serde_json::to_value(document.with_source(replacement.into())).unwrap(),
expected
);
}
#[rstest::rstest]
#[case::document_url("document_url", "document_name", "application/pdf")]
#[case::image_url("image_url", "detail", "image/png")]
fn document_variants_preserve_provider_fields_when_rewriting_sources(
#[case] kind: &str,
#[case] field: &str,
#[case] mime_type: &str,
#[values(json!("kept"), Value::Null)] extra: Value,
) {
let original = "https://example.com/input";
let replacement = format!("data:{mime_type};base64,AA==");
let document: OcrDocument =
serde_json::from_value(json!({"type": kind, kind: original, field: extra})).unwrap();
assert_eq!(document.source(), original);
assert_eq!(
serde_json::to_value(document.with_source(replacement.clone())).unwrap(),
json!({"type": kind, kind: replacement, field: extra})
);
}
#[test]

View file

@ -956,32 +956,29 @@ async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_over
}
}
#[rstest]
#[case::declared("Content-Length: 1000000")]
#[case::chunked("Transfer-Encoding: chunked")]
#[tokio::test]
async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() {
let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1));
for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] {
let body = if headers.starts_with("Transfer") {
format!("{:x}\r\n{prefix}\r\n", prefix.len())
} else {
prefix.clone()
};
let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}");
let error = read_bounded_response(response.into_bytes(), 4096)
.await
.unwrap_err();
match error {
super::Error::Transport(crate::transport::Error::Http { status, body }) => {
assert_eq!(status, 429);
assert_eq!(
body,
format!(
"{}... (truncated)",
"x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS)
)
);
}
error => panic!("unexpected error: {error}"),
async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining(
#[case] headers: &str,
) {
let prefix = "x".repeat(4096);
let body = if headers.starts_with("Transfer") {
format!("{:x}\r\n{prefix}\r\n", prefix.len())
} else {
prefix.clone()
};
let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}");
let error = read_bounded_response(response.into_bytes(), prefix.len())
.await
.unwrap_err();
match error {
super::Error::Transport(crate::transport::Error::Http { status, body }) => {
assert_eq!(status, 429);
assert_eq!(body, prefix);
}
error => panic!("unexpected error: {error}"),
}
}

View file

@ -1,25 +1,59 @@
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};
pub(super) fn to_pyerr(error: Error) -> PyErr {
let status = error.http_status_code();
let mapped = match error {
Error::Provider { status, body, .. }
| Error::Transport(litellm_core::transport::Error::Http { status, body }) => {
RustUpstreamError::new_err((status, body))
}
Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => {
PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
}
Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()),
other => core_error_to_pyerr(other.into()),
};
let mapped = Python::attach(|py| -> PyResult<PyErr> {
Ok(match error {
Error::Provider {
status,
body,
headers,
} => upstream_error(py, status, body, headers)?,
Error::Transport(litellm_core::transport::Error::Http { status, body }) => {
upstream_error(py, status, body, Vec::new())?
}
Error::RequestFormat => {
let error = core_error_to_pyerr(Error::RequestFormat.into());
error
.value(py)
.setattr("ocr_request_format_error", true)
.ok();
error
}
Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => {
PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
}
Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()),
other => core_error_to_pyerr(other.into()),
})
})
.unwrap_or_else(|error| error);
attach_status(mapped, status)
}
fn upstream_error(
py: Python<'_>,
status: u16,
body: String,
headers: Vec<(String, String)>,
) -> PyResult<PyErr> {
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)?;
Ok(error)
}
fn attach_status(error: PyErr, status: Option<u16>) -> PyErr {
if let Some(status) = status {
Python::attach(|py| {
@ -50,7 +84,7 @@ mod tests {
.unwrap()
.extract::<u16>()
.unwrap(),
500
400
);
let mapped = to_pyerr(Error::Provider {
status: 429,

View file

@ -88,7 +88,21 @@ enum ProjectedDocument {
impl ProjectedDocument {
fn project(document: &Bound<'_, PyAny>) -> PyResult<Self> {
let kind: String = document.get_item("type")?.extract()?;
let kind: String = document
.get_item("type")
.and_then(|value| value.extract())
.map_err(|error| {
let py = document.py();
if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py)
|| error.is_instance_of::<pyo3::exceptions::PyTypeError>(py)
{
ocr_error_to_pyerr(litellm_core::ocr::Error::RequestField {
path: "document.type".into(),
})
} else {
error
}
})?;
if kind != "file" {
return Ok(Self::Other {
wire: from_py(document)?,
@ -185,7 +199,7 @@ pub(super) fn admitted_call(outcome: NativeOutcome<OcrCall>) -> PyResult<OcrCall
mod tests {
use litellm_core::ocr::Error;
use litellm_core::ocr::OcrDecline;
use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
use pyo3::exceptions::PyValueError;
use super::*;
@ -515,21 +529,21 @@ kwargs = {'api_key': key}
}
#[test]
fn document_discriminator_errors_keep_their_existing_exceptions() {
fn document_discriminator_errors_are_validation_errors_and_preserve_custom_failures() {
Python::initialize();
Python::attach(|py| {
let missing = py.eval(c"{}", None, None).unwrap();
assert!(
project_document(&missing)
.unwrap_err()
.is_instance_of::<PyKeyError>(py)
.is_instance_of::<PyValueError>(py)
);
let non_string = py.eval(c"{'type': 1}", None, None).unwrap();
assert!(
project_document(&non_string)
.unwrap_err()
.is_instance_of::<PyTypeError>(py)
.is_instance_of::<PyValueError>(py)
);
let locals = eval(

View file

@ -500,6 +500,7 @@ class RateLimitError(openai.RateLimitError):
self.response = httpx.Response(
status_code=429,
headers=_response_headers,
content=response.content if response is not None else None,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",

View file

@ -59,7 +59,7 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
from litellm.llms.base_llm.ocr.transformation import OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, OCRResponse
from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
@ -1568,7 +1568,7 @@ class BaseLLMHTTPHandler:
transformed_result: Final = provider_config.transform_ocr_request(
model=model,
document=document,
optional_params=optional_params,
optional_params={key: value for key, value in optional_params.items() if key != OCR_REQUEST_FORMAT_PARAM},
headers=headers,
api_key=api_key,
api_base=api_base,
@ -1634,7 +1634,7 @@ class BaseLLMHTTPHandler:
transformed_result: Final = await provider_config.async_transform_ocr_request(
model=model,
document=document,
optional_params=optional_params,
optional_params={key: value for key, value in optional_params.items() if key != OCR_REQUEST_FORMAT_PARAM},
headers=headers,
api_key=api_key,
api_base=api_base,
@ -1672,12 +1672,26 @@ class BaseLLMHTTPHandler:
optional_params: Mapping[str, object],
) -> OCRResponse:
"""Shared logic for transforming OCR responses."""
return provider_config.transform_ocr_response(
normalized: Final = provider_config.transform_ocr_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
optional_params=optional_params,
)
return self._finalize_ocr_response(normalized, response, optional_params)
@staticmethod
def _finalize_ocr_response(
normalized: OCRResponse,
response: httpx.Response,
optional_params: Mapping[str, object],
) -> OCRResponse:
if (
optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native"
and normalized.get_provider_native_response() is None
):
normalized.set_provider_native_response(response.json())
return normalized
def ocr(
self,
@ -1823,12 +1837,13 @@ class BaseLLMHTTPHandler:
)
# Use async response transform for async operations
return await provider_config.async_transform_ocr_response(
normalized: Final = await provider_config.async_transform_ocr_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
optional_params=optional_params,
)
return self._finalize_ocr_response(normalized, response, optional_params)
def search(
self,
@ -6157,6 +6172,8 @@ class BaseLLMHTTPHandler:
status_code=status_code,
headers=error_headers,
)
if isinstance(provider_config, BaseOCRConfig) and isinstance(error_response, httpx.Response):
provider_error.response = error_response
if not isinstance(received_status_code, int):
provider_error.status_code_is_synthesized = True
raise provider_error

View file

@ -70,16 +70,27 @@ def _prepare_ocr_request(
)
if not isinstance(document, dict):
raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}")
raise litellm.BadRequestError(
message="document must be a dict with 'type' and URL/file field",
model=model,
llm_provider=custom_llm_provider or "",
)
doc_type = document.get("type")
normalized_document: Final = (
convert_file_document_to_url_document(document) if document.get("type") == "file" else document
)
doc_type: Final = normalized_document.get("type")
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'")
if doc_type not in ("document_url", "image_url"):
raise litellm.BadRequestError(
message=f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'",
model=model,
llm_provider=custom_llm_provider or "",
)
if not normalized_document.get(doc_type):
raise litellm.BadRequestError(
message="Document URL is required", model=model, llm_provider=custom_llm_provider or ""
)
(
model,
@ -116,31 +127,26 @@ def _prepare_ocr_request(
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
non_default_params: Final = {param: kwargs.pop(param) for param in supported_params if param in kwargs}
optional_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
try:
mapped_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
except ValueError as error:
raise litellm.BadRequestError(message=str(error), model=model, llm_provider=custom_llm_provider) from error
optional_params: Final = {
**mapped_params,
**({OCR_REQUEST_FORMAT_PARAM: requested_format} if requested_format is not None else {}),
}
verbose_logger.debug("OCR optional_params after mapping: %s", optional_params)
@ -160,7 +166,7 @@ def _prepare_ocr_request(
return _PreparedOCRRequest(
model=model,
document=document,
document=normalized_document,
api_key=resolved_api_key,
api_base=resolved_api_base,
custom_llm_provider=custom_llm_provider,

View file

@ -3,6 +3,8 @@ from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
import httpx
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
@ -51,17 +53,28 @@ def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]:
def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception:
model: Final = request.model.removeprefix(f"{request_provider}/")
if getattr(error, "ocr_request_format_error", False):
return litellm.UnsupportedParamsError(
message=f"Invalid `req_format`: {request.kwargs.get('req_format')!r}. Expected 'native' or 'litellm'.",
model=model,
llm_provider=request_provider,
)
mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper
ExceptionMapper, litellm.exception_type
)
try:
return mapper(
model=request.model.removeprefix(f"{request_provider}/"),
model=model,
custom_llm_provider=request_provider,
original_exception=error,
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
public_error.__context__ = error
return public_error