From 5cb55a5efdcca8ff156f0eb36146567c598fe96f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 7 Sep 2026 08:05:57 -0700 Subject: [PATCH] wip --- .github/workflows/test-rust.yml | 17 + Makefile | 20 +- litellm-rust/README.md | 23 +- litellm-rust/crates/core/src/constants.rs | 2 + litellm-rust/crates/core/src/http_utils.rs | 2 + .../core/src/http_utils/buffered_post.rs | 74 ++ .../crates/core/tests/buffered_post.rs | 45 + .../crates/python-bridge/src/errors.rs | 38 - .../crates/python-bridge/src/execution.rs | 37 +- .../python-bridge/src/routes/definition.rs | 8 +- .../crates/python-bridge/src/routes/ocr.rs | 236 +++-- .../tests/fixtures/callback_components.py | 20 +- .../tests/fixtures/callback_lifecycle.py | 4 + .../python-interop/tests/prepared_call.rs | 65 +- litellm/ocr/main.py | 197 +--- litellm/rust_bridge/ocr.py | 197 ++-- ...cr_azure_document_intelligence_api_base.py | 45 +- tests/test_litellm/ocr/test_rust_bridge.py | 990 ++++++++++++------ .../rust_bridge/native_route_wheel_test.py | 69 +- 19 files changed, 1373 insertions(+), 716 deletions(-) create mode 100644 litellm-rust/crates/core/src/http_utils/buffered_post.rs create mode 100644 litellm-rust/crates/core/tests/buffered_post.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index aaf31514679..ee1bb748342 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,13 @@ on: push: paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -21,6 +28,13 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - "litellm/rust_bridge/**" + - "litellm/ocr/**" + - "litellm/llms/base_llm/ocr/**" + - "litellm/llms/custom_httpx/llm_http_handler.py" + - "tests/test_litellm/ocr/**" + - "tests/test_litellm/conftest.py" + - "Makefile" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -128,6 +142,9 @@ jobs: - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl + - name: Require native OCR pytest acceptance + run: make test-rust-ocr RUST_OCR_WHEEL="$(realpath dist/*.whl)" + - name: Check Python fixtures for Cargo tests run: make lint-rust-python-fixtures diff --git a/Makefile b/Makefile index 0fd9d8032b0..fcffdb02e07 100644 --- a/Makefile +++ b/Makefile @@ -11,7 +11,7 @@ lint-test-quality lint-test-quality-budget-update \ install-dev install-proxy-dev install-test-deps install-hooks \ install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \ - lint-install lint-fetch-base bootstrap install-rust-python-test-deps test-rust-python lint-rust-python-fixtures + lint-install lint-fetch-base bootstrap install-rust-python-test-deps test-rust-python test-rust-ocr lint-rust-python-fixtures # Default target help: @@ -57,6 +57,7 @@ help: @echo " make test-unit-helm - Run helm unit tests" @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @echo " make test-rust-python - Run ignored Python-integrated Cargo tests" + @echo " make test-rust-ocr - Build a wheel and require native OCR pytest acceptance" @echo " make lint-rust-python-fixtures - Check Rust test Python fixtures with Ruff" @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @@ -307,6 +308,23 @@ test-rust-extension: LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust +test-rust-ocr: + @temporary=$$(mktemp -d) && \ + trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \ + if [ -n "$(RUST_OCR_WHEEL)" ]; then \ + wheel="$(RUST_OCR_WHEEL)"; \ + else \ + $(UV) build --wheel --out-dir "$$temporary/wheels" || exit $$?; \ + set -- "$$temporary"/wheels/*.whl; \ + [ "$$#" -eq 1 ] || exit 1; \ + wheel="$$1"; \ + fi && \ + UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --no-default-groups --group dev --extra proxy && \ + $(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$wheel" && \ + LITELLM_LOCAL_MODEL_COST_MAP=True "$$temporary/venv/bin/python" -I -c 'from litellm.rust_bridge import _native; assert callable(_native.ocr) and callable(_native.aocr)' && \ + LITELLM_REQUIRE_NATIVE_OCR=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ + "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib tests/test_litellm/ocr/test_rust_bridge.py -v + test-rust-python: install-rust-python-test-deps @python=$$($(UV_RUN) python -c 'import sys; print(sys.executable)') && \ site_packages=$$("$$python" -c 'import os, sysconfig; print(os.pathsep.join(dict.fromkeys(sysconfig.get_path(key) for key in ("purelib", "platlib"))))') && \ diff --git a/litellm-rust/README.md b/litellm-rust/README.md index cf82418c0b0..68e02ab335d 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -73,7 +73,28 @@ make lint-rust-python-fixtures `lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without syncing the project environment -These tests validate retained callback identity, mutation, invocation context, +Run the native OCR acceptance gate from the repository root: + +```bash +make test-rust-ocr +``` + +This builds the current release wheel, installs locked SDK dependencies, the +`dev` test group, and the `proxy` extra in a temporary Python 3.12 environment, +then installs the wheel without resolving dependencies again. The proxy extra +is needed by the shared pytest fixtures. Python isolated mode and pytest's +importlib mode keep the checkout from shadowing the installed wheel + +The gate checks that native `ocr` and `aocr` are importable, then runs +`tests/test_litellm/ocr/test_rust_bridge.py` with +`LITELLM_REQUIRE_NATIVE_OCR=1`, so unavailable native OCR fails instead of +skipping. CI uses `make test-rust-ocr RUST_OCR_WHEEL=/absolute/path/to/current.whl` +to test the release wheel it just built. The stdlib-only +`native_route_wheel_test.py` also exercises sync/async OCR through a small +boundary, including 429 handling in `finish`/`afinish`, alongside the other +native routes + +The Python-integrated Cargo tests validate retained callback identity, mutation, invocation context, and ownership against Python behavior, including existing LiteLLM components. They do not wire retained callbacks into production routes or change provider preparation, authentication, HTTP transport, or response transformation diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index fc81f4fa029..ec1a2988e75 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -32,6 +32,8 @@ pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; +pub(crate) const BUFFERED_POST_CONNECT_TIMEOUT_SECS: u64 = 10; + /// `object` field every non-streaming chat completion response carries. pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index cb472dd5a57..93dfc1a5a8f 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -1,5 +1,7 @@ //! Header and upstream-body helpers shared by every route module. +pub mod buffered_post; + use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; diff --git a/litellm-rust/crates/core/src/http_utils/buffered_post.rs b/litellm-rust/crates/core/src/http_utils/buffered_post.rs new file mode 100644 index 00000000000..28deafb3517 --- /dev/null +++ b/litellm-rust/crates/core/src/http_utils/buffered_post.rs @@ -0,0 +1,74 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; + +use crate::constants::BUFFERED_POST_CONNECT_TIMEOUT_SECS; +use crate::error::Error; + +pub struct Request { + pub url: String, + pub headers: Vec<(Vec, Vec)>, + pub body: Vec, + pub timeout_seconds: f64, +} + +pub struct Response { + pub status: u16, + pub headers: Vec<(Vec, Vec)>, + pub content: Vec, +} + +pub async fn send(request: Request) -> Result { + let timeout = Duration::try_from_secs_f64(request.timeout_seconds) + .ok() + .filter(|timeout| !timeout.is_zero()) + .ok_or_else(|| Error::InvalidRequest("timeout must be positive and finite".into()))?; + let mut headers = HeaderMap::new(); + for (name, value) in request.headers { + let name = HeaderName::from_bytes(&name) + .map_err(|_| Error::InvalidRequest("invalid header name".into()))?; + let value = HeaderValue::from_bytes(&value) + .map_err(|_| Error::InvalidRequest("invalid header value".into()))?; + headers.append(name, value); + } + + static CLIENT: OnceLock> = OnceLock::new(); + let client = CLIENT + .get_or_init(|| { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(BUFFERED_POST_CONNECT_TIMEOUT_SECS)) + .redirect(reqwest::redirect::Policy::none()) + .no_gzip() + .no_brotli() + .no_deflate() + .no_zstd() + .build() + }) + .as_ref() + .map_err(|_| Error::Network("could not initialize HTTP client".into()))?; + let response = client + .post(request.url) + .headers(headers) + .body(request.body) + .timeout(timeout) + .send() + .await + .map_err(|_| Error::Network("transport failed".into()))?; + let status = response.status().as_u16(); + let headers = response + .headers() + .iter() + .map(|(name, value)| (name.as_str().as_bytes().to_vec(), value.as_bytes().to_vec())) + .collect(); + let content = response + .bytes() + .await + .map_err(|_| Error::Network("could not read response".into()))? + .to_vec(); + Ok(Response { + status, + headers, + content, + }) +} diff --git a/litellm-rust/crates/core/tests/buffered_post.rs b/litellm-rust/crates/core/tests/buffered_post.rs new file mode 100644 index 00000000000..c79a2ec8d25 --- /dev/null +++ b/litellm-rust/crates/core/tests/buffered_post.rs @@ -0,0 +1,45 @@ +use litellm_core::error::Error; +use litellm_core::http_utils::buffered_post::{Request, send}; + +fn request() -> Request { + Request { + url: "unknown://private-document?secret=credential".into(), + headers: vec![], + body: vec![0, 255], + timeout_seconds: 1.0, + } +} + +#[tokio::test] +async fn rejects_invalid_timeouts_and_headers_without_echoing_wire_values() { + for timeout_seconds in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::MAX] { + let result = send(Request { + timeout_seconds, + ..request() + }) + .await; + assert!( + matches!(result, Err(Error::InvalidRequest(message)) if message == "timeout must be positive and finite") + ); + } + for (headers, expected) in [ + ( + vec![(b"private\nname".to_vec(), b"secret".to_vec())], + "invalid header name", + ), + ( + vec![(b"x-proof".to_vec(), b"private\nvalue".to_vec())], + "invalid header value", + ), + ] { + let result = send(Request { + headers, + ..request() + }) + .await; + assert!(matches!(result, Err(Error::InvalidRequest(message)) if message == expected)); + } + assert!( + matches!(send(request()).await, Err(Error::Network(message)) if message == "transport failed") + ); +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 76c298abf89..914e2e1e033 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -59,41 +59,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } - -pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::MissingField("document_url" | "image_url") => { - PyValueError::new_err("Document URL is required") - } - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - other => core_error_to_pyerr(other), - } -} - -#[cfg(test)] -mod ocr_error_tests { - use super::*; - - #[test] - fn ocr_errors_preserve_python_validation_and_provider_details() { - Python::initialize(); - Python::attach(|py| { - for field in ["document_url", "image_url"] { - let mapped = ocr_error_to_pyerr(Error::MissingField(field)); - assert!(mapped.is_instance_of::(py)); - assert_eq!(mapped.value(py).to_string(), "Document URL is required"); - } - let mapped = ocr_error_to_pyerr(Error::Http { - status: 429, - body: r#"{"message":"rate limited"}"#.to_string(), - }); - assert!(mapped.is_instance_of::(py)); - let args: (u16, String) = mapped - .value(py) - .getattr("args") - .and_then(|args| args.extract()) - .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index f3648158cf6..ca0f8a4ae86 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -20,7 +20,20 @@ where T: Serialize + Send + 'static, F: Future> + Send + 'static, { - run_sync_on( + let result = run_sync_value(py, future, map_error)?; + Pythonized(result).into_pyobject(py).map(Bound::unbind) +} + +pub(crate) fn run_sync_value( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_value_on( py, pyo3_async_runtimes::tokio::get_runtime(), future, @@ -28,14 +41,14 @@ where ) } -fn run_sync_on( +fn run_sync_value_on( py: Python<'_>, runtime: &Runtime, future: F, map_error: fn(Error) -> PyErr, -) -> PyResult> +) -> PyResult where - T: Serialize + Send + 'static, + T: Send + 'static, F: Future> + Send + 'static, { if Handle::try_current().is_ok() { @@ -45,8 +58,7 @@ where } let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; - let result = map_core_result(result, map_error)?; - Pythonized(result).into_pyobject(py).map(Bound::unbind) + map_core_result(result, map_error) } pub(crate) fn run_async( @@ -65,6 +77,15 @@ where }) } +pub(crate) async fn run_async_value(future: F, map_error: fn(Error) -> PyErr) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + let result = catch_future_panic(future).await?; + map_core_result(result, map_error) +} + fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), @@ -256,7 +277,7 @@ mod tests { .build() .expect("runtime should build"); Python::attach(|py| { - let result = run_sync_on( + let result = run_sync_value_on( py, &runtime, async { @@ -265,7 +286,7 @@ mod tests { }, runtime_error, ); - assert!(extract_bool(py, result)); + assert!(result.expect("route should complete")); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index bc51647cbad..843c338e8db 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -222,11 +222,7 @@ mod tests { let module = PyModule::new(py, "routes").expect("module should be created"); crate::routes::register(&module).expect("routes should register"); let routes = [ - ( - "ocr", - "aocr", - "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", - ), + ("ocr", "aocr", "(boundary)"), ( "transcription", "atranscription", @@ -311,7 +307,7 @@ mod tests { .expect("kwargs should accept extra_headers"); let document = PyDict::new(py); - for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] { + for (sync_name, async_name) in [("transcription", "atranscription")] { let sync_error = module .getattr(sync_name) .and_then(|function| function.call(("model", &document), Some(&kwargs))) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index cc2f8e43cea..172b8de48a7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -1,73 +1,181 @@ -use litellm_core::Error; -use std::future::Future; +//! Retained OCR route: Python owns request/response objects, Rust sequences +//! prepare -> encode -> POST -> finish through owning `Py` handles. -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; +use litellm_core::http_utils::buffered_post::{self, Request, Response}; +use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; use pyo3::prelude::*; -use serde_json::Value; +use pyo3::sync::PyOnceLock; +use pyo3::types::{PyBytes, PyList, PyTuple}; -use crate::errors::ocr_error_to_pyerr; -use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; +use crate::errors::core_error_to_pyerr; +use crate::execution::{run_async_value, run_sync_value}; -fn prepare_ocr( - inputs: OcrInputs, -) -> PyResult> + Send + 'static> { - let document = inputs.document; - let options = RouteOptions::from_python(RouteOptionsInputs { - model: inputs.model, - api_key: inputs.api_key, - api_base: inputs.api_base, - custom_llm_provider: inputs.custom_llm_provider, - extra_headers: inputs.extra_headers, - timeout_seconds: inputs.timeout_seconds, - })?; - let optional_params = object_or_empty("optional_params", inputs.optional_params)?; +#[derive(Clone, Copy)] +struct BoundaryStep { + method: &'static str, + awaited: bool, +} - Ok(async move { - let RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout, - } = options; - run_ocr(OcrRequest { - model: &model, - document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await +const PREPARE_SYNC: BoundaryStep = BoundaryStep { + method: "prepare", + awaited: false, +}; +const PREPARE_ASYNC: BoundaryStep = BoundaryStep { + method: "aprepare", + awaited: true, +}; +const ENCODE: BoundaryStep = BoundaryStep { + method: "encode", + awaited: false, +}; +const FINISH_SYNC: BoundaryStep = BoundaryStep { + method: "finish", + awaited: false, +}; +const FINISH_ASYNC: BoundaryStep = BoundaryStep { + method: "afinish", + awaited: true, +}; + +fn invoke( + boundary: &Bound<'_, PyAny>, + step: BoundaryStep, + args: Bound<'_, PyTuple>, +) -> PyResult> { + let call = PreparedCall::new( + if step.awaited { + InvocationMode::Await + } else { + InvocationMode::Direct + }, + boundary.getattr(step.method)?.unbind(), + args.unbind(), + None, + ); + match call.invoke(boundary.py())? { + InvocationOutcome::Returned(value) | InvocationOutcome::Awaitable(value) => Ok(value), + } +} + +#[pyfunction] +fn prepare(boundary: &Bound<'_, PyAny>, asynchronous: bool) -> PyResult> { + let step = if asynchronous { + PREPARE_ASYNC + } else { + PREPARE_SYNC + }; + invoke(boundary, step, PyTuple::empty(boundary.py())) +} + +fn request(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult { + type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; + let py = boundary.py(); + let encoded = invoke(boundary, ENCODE, PyTuple::new(py, [roots])?)?; + let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) = + encoded.into_bound(py).extract()?; + Ok(Request { + url, + headers: headers + .into_iter() + .map(|(name, value)| (name.as_bytes().to_vec(), value.as_bytes().to_vec())) + .collect(), + body: body.as_bytes().to_vec(), + timeout_seconds, }) } -bridge_route! { - sync = ocr, - asynchronous = aocr, - inputs = OcrInputs, - required = { - model: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - document: serde_json::Value, - }, - optional = { - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, - timeout_seconds: Option, - }, - prepare = prepare_ocr, - errors = ocr_error_to_pyerr, +struct Wire(Response); + +impl<'py> IntoPyObject<'py> for Wire { + type Target = PyTuple; + type Output = Bound<'py, PyTuple>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + let headers = PyList::new( + py, + self.0 + .headers + .iter() + .map(|(name, value)| (PyBytes::new(py, name), PyBytes::new(py, value))), + )?; + (self.0.status, headers, PyBytes::new(py, &self.0.content)).into_pyobject(py) + } +} + +#[pyfunction] +fn send<'a>( + boundary: &'a Bound<'a, PyAny>, + roots: &Bound<'_, PyAny>, +) -> PyResult> { + let request = request(boundary, roots)?; + pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move { + let response = run_async_value(buffered_post::send(request), core_error_to_pyerr).await?; + Ok(Wire(response)) + }) +} + +#[pyfunction] +fn finish( + boundary: &Bound<'_, PyAny>, + wire: &Bound<'_, PyAny>, + asynchronous: bool, +) -> PyResult> { + let step = if asynchronous { + FINISH_ASYNC + } else { + FINISH_SYNC + }; + invoke(boundary, step, PyTuple::new(boundary.py(), [wire])?) +} + +#[pyfunction] +fn ocr(boundary: &Bound<'_, PyAny>) -> PyResult> { + let py = boundary.py(); + let roots = prepare(boundary, false)?; + let request = request(boundary, roots.bind(py))?; + let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?; + let wire = Wire(response).into_pyobject(py)?; + finish(boundary, &wire, false) +} + +#[pyfunction] +fn aocr<'a>(boundary: &'a Bound<'a, PyAny>) -> PyResult> { + driver(boundary.py())?.getattr("drive")?.call1((boundary,)) +} + +/// The async route must await `aprepare`/`afinish` inline in the caller's +/// Python task, so a Python driver coroutine owns the roots between steps. +/// Compiling the driver runs Python (audit hooks can re-enter `aocr`), so +/// compile first and publish only a finished module into the once-lock. +fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { + static DRIVER: PyOnceLock> = PyOnceLock::new(); + if let Some(module) = DRIVER.get(py) { + return Ok(module.bind(py)); + } + let module = PyModule::from_code( + py, + c"async def drive(boundary): + roots = await _prepare(boundary, True) + wire = await _send(boundary, roots) + return await _finish(boundary, wire, True) +", + c"ocr_driver.py", + c"_ocr_driver", + )?; + module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; + module.add("_send", wrap_pyfunction!(send, &module)?)?; + module.add("_finish", wrap_pyfunction!(finish, &module)?)?; + Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!(ocr, module)?)?; + crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!(aocr, module)?)?; + Ok(()) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + register(module) } diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py index 6789d962c53..d5379a9dc98 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py @@ -35,6 +35,8 @@ def logger_for(callbacks=(), stream=False, input_callbacks=(), sync_callbacks=() async def real_pre_call_logging(owners): retained = [] observed = [] + snapshots = [] + order = [] ignored = {"replacement": True} metadata = {"secret": "private", "keep": []} removed = object() @@ -42,11 +44,13 @@ async def real_pre_call_logging(owners): class Retain(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append("retain") retained.append(kwargs) return ignored class Mutate(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append("mutate") kwargs["normalized"] = "normalized" assert kwargs.pop("remove") is removed kwargs["retained_metadata"]["secret"] = "masked" @@ -54,12 +58,24 @@ async def real_pre_call_logging(owners): class Fail(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append("fail") kwargs["lock"] = lock kwargs["retained_metadata"]["keep"].append("before failure") raise RuntimeError("expected pre-call callback failure") class Observe(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append("observe") + snapshots.append( + ( + kwargs.get("normalized"), + "remove" in kwargs, + kwargs["retained_metadata"]["secret"], + tuple(kwargs["retained_metadata"]["keep"]), + "lock" in kwargs, + "replacement" in kwargs, + ) + ) observed.append((kwargs, messages)) logger = logger_for(input_callbacks=[Retain(), Mutate(), Fail(), Observe()]) @@ -72,7 +88,9 @@ async def real_pre_call_logging(owners): assert owner.invoke() is None finally: owner.close() - assert retained == [details] and observed == [(details, messages)] + assert order == ["retain", "mutate", "fail", "observe"] + assert snapshots == [("normalized", False, "masked", ("before failure",), True, False)] + assert len(retained) == len(observed) == 1 assert retained[0] is details and observed[0][0] is details assert observed[0][1] is messages and details["input"] is messages assert details["additional_args"] is additional diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py index c1e133db971..f046595158c 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py @@ -661,15 +661,19 @@ async def detached_work_after_error(owners): def run_checked(owners, scenario): baseline = owners.live + background_failures = [] async def run(): + asyncio.get_running_loop().set_exception_handler(lambda loop, context: background_failures.append(context)) await asyncio.wait_for(scenario, timeout=15) + gc.collect() assert owners.live == baseline pending = asyncio.all_tasks() - {asyncio.current_task()} assert not pending, f"undrained tasks: {pending}" asyncio.run(run()) gc.collect() + assert not background_failures, f"unhandled background failures: {background_failures}" assert owners.live == baseline diff --git a/litellm-rust/crates/python-interop/tests/prepared_call.rs b/litellm-rust/crates/python-interop/tests/prepared_call.rs index a8549720c9d..d5d444d31b6 100644 --- a/litellm-rust/crates/python-interop/tests/prepared_call.rs +++ b/litellm-rust/crates/python-interop/tests/prepared_call.rs @@ -6,7 +6,7 @@ use rstest::rstest; #[path = "support/mod.rs"] mod support; -use support::python::{InitializedPython, initialized_python, item, scope}; +use support::python::{InitializedPython, initialized_python, item, run_fixture, scope}; #[rstest] fn retains_aliases_mutations_and_original_result( @@ -271,6 +271,55 @@ fn prepare_pre_call( )) } +#[rstest] +fn checked_runner_rejects_unhandled_background_failures( + initialized_python: &InitializedPython, +) -> PyResult<()> { + let _ = initialized_python; + Python::attach(|py| { + let globals = PyDict::new(py); + run_fixture( + py, + &globals, + include_str!("fixtures/callback_lifecycle.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/callback_lifecycle.py" + ), + )?; + py.run( + c" +async def fail(): + raise RuntimeError('background task regression') + +for cyclic in (False, True): + for handled in (False, True): + async def scenario(cyclic=cyclic, handled=handled): + task = asyncio.create_task(fail()) + if cyclic: + task.cycle = task + await checkpoint() + assert task.done() + if handled: + with TestCase().assertRaisesRegex(RuntimeError, 'background task regression'): + task.result() + del task + + owners = ReferenceFactory() + if handled: + run_checked(owners, scenario()) + else: + with TestCase().assertRaisesRegex( + AssertionError, r'unhandled background failures: .*background task regression' + ): + run_checked(owners, scenario()) +", + Some(&globals), + None, + ) + }) +} + #[rstest] #[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] fn real_ocr_logging_preserves_execution_roots_and_continues_after_error( @@ -285,15 +334,20 @@ from datetime import datetime from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging +order = [] + class Retain(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append('retain') self.view = kwargs['additional_args'] self.headers = self.view['headers'] self.body = self.view['complete_input_dict'] + self.snapshot = (self.headers['X-Trace'], self.body['document']['value']) return {'ignored_replacement': True} class MutateThenFail(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append('mutate_then_fail') view = kwargs['additional_args'] view['headers']['X-Trace'] = 'mutated' view['complete_input_dict']['document']['value'] = 'mutated' @@ -303,7 +357,13 @@ class MutateThenFail(CustomLogger): class Observe(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): + order.append('observe') self.view = kwargs['additional_args'] + self.snapshot = ( + tuple(sorted(self.view['headers'].items())), + self.view['complete_input_dict'].get('replacement'), + 'document' in self.view['complete_input_dict'], + ) first = Retain() last = Observe() @@ -338,6 +398,9 @@ logger = Logging( ); py.run( c" +assert order == ['retain', 'mutate_then_fail', 'observe'] +assert first.snapshot == ('original', 'original') +assert last.snapshot == ((('X-Trace', 'replacement'),), True, False) assert first.view is last.view assert first.body['document'] is document assert first.body['alias'] is document diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 74f192b0115..0164d829bec 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -7,7 +7,7 @@ import base64 import mimetypes import os import re -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import Coroutine, Mapping from dataclasses import dataclass from io import IOBase from typing import Any, Final, cast @@ -29,7 +29,6 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -54,14 +53,6 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - _RUST_OCR_PROVIDERS: Final = { "mistral", "azure_ai", @@ -198,169 +189,39 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS -def _rust_bridge_optional_params( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> dict[str, object]: - optional_params: Final = dict(prepared_request.optional_params) - if prepared_request.custom_llm_provider == "vertex_ai": - vertex_project: Final = ( - prepared_request.litellm_params.get("vertex_project") - or prepared_request.litellm_params.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - vertex_location: Final = ( - prepared_request.litellm_params.get("vertex_location") - or prepared_request.litellm_params.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - if vertex_project is not None: - optional_params["vertex_project"] = vertex_project - if vertex_location is not None: - optional_params["vertex_location"] = vertex_location - return optional_params +def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: + raw_request_override: Final = prepared_request.litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + return rust_enabled(request_override=request_override) -def _rust_bridge_api_base( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> str | None: - if prepared_request.api_base is not None: - return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai": - if is_azure_document_intelligence_model(prepared_request.model): - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - return resolve_secret("AZURE_AI_API_BASE") - return None - - -def _prepare_rust_ocr_call( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: - provider_config: Final = prepared_request.provider_config - api_key_env_var: Final = provider_config.get_api_key_env_var() - resolved_api_key: Final = prepared_request.api_key or ( - resolve_api_key(api_key_env_var) if api_key_env_var is not None else None - ) - resolved_headers: Final = provider_config.validate_environment( - headers=prepared_request.extra_headers or {}, - model=prepared_request.model, - api_key=resolved_api_key, - api_base=prepared_request.api_base, - litellm_params=prepared_request.litellm_params, - ) - resolved_complete_url: Final = provider_config.get_complete_url( - api_base=prepared_request.api_base, +def _ocr_boundary(prepared_request: _PreparedOCRRequest) -> rust_ocr_bridge.OCRBoundary: + return rust_ocr_bridge.OCRBoundary( + handler=base_llm_http_handler, model=prepared_request.model, + document=prepared_request.document, optional_params=prepared_request.optional_params, + logging_obj=prepared_request.litellm_logging_obj, + api_key=prepared_request.api_key, + api_base=prepared_request.api_base, + headers=prepared_request.extra_headers, + provider_config=prepared_request.provider_config, litellm_params=prepared_request.litellm_params, - ) - rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) - rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) - prepared_request.litellm_logging_obj.pre_call( - input="OCR document processing", - api_key=resolved_api_key, - additional_args={ - "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, - "api_base": resolved_complete_url, - "headers": resolved_headers, - }, - ) - return _PreparedRustOCRCall( - api_key=resolved_api_key, - api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, + custom_llm_provider=prepared_request.custom_llm_provider, + timeout=prepared_request.effective_timeout, ) -def _map_rust_ocr_error( - error: Exception, - prepared_request: _PreparedOCRRequest, - exception_types: tuple[type[BaseException], type[BaseException]] | None, -) -> Exception: - if exception_types is None: - return error - _, upstream_error = exception_types - if not isinstance(error, upstream_error): - return error - error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs - tuple[object, ...], error.args - ) - status_value: Final = error_args[0] if error_args else 0 - message_value: Final = error_args[1] if len(error_args) > 1 else str(error) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) - error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped - Callable[..., Exception], prepared_request.provider_config.get_error_class - ) - return error_factory( - error_message=message, - status_code=status or 500, - headers={}, # mutable-ok: provider error factories require a concrete header dict - ) -def _run_rust_ocr( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: +def _run_rust_ocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None: if rust_ocr_bridge.load_rust_ocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - try: - rust_response: Final = rust_ocr_bridge.ocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + return rust_ocr_bridge.ocr(_ocr_boundary(prepared_request)) -async def _run_rust_aocr( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: +async def _run_rust_aocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None: if rust_ocr_bridge.load_rust_aocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) - try: - rust_response: Final = await rust_ocr_bridge.aocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + return await rust_ocr_bridge.aocr(_ocr_boundary(prepared_request)) @client @@ -457,13 +318,8 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + rust_response: Final = await _run_rust_aocr(prepared_request=prepared) if rust_response is None: verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") else: @@ -729,13 +585,8 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): + rust_response: Final = _run_rust_ocr(prepared_request=prepared) if rust_response is None: verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") else: diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index b7fdb5a98ef..2a412957119 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -1,44 +1,145 @@ -"""Thin Python wrapper for the native Rust OCR bridge.""" +"""Retained OCR bridge: Python owns the request/response objects, Rust drives +prepare -> encode -> POST -> finish against those same objects.""" from __future__ import annotations +import math from collections.abc import Awaitable +from dataclasses import dataclass, field from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, DocumentType, OCRResponse +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, + get_async_httpx_client, +) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.rust_bridge import configuration as _configuration from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds +rust_ocr_enabled = _configuration.rust_enabled +rust = _configuration.rust + +OCRRoots = tuple[dict[str, object], str, dict[str, object], None] +OCRWire = tuple[int, list[tuple[bytes, bytes]], bytes] +OCREncoded = tuple[str, list[tuple[bytes, bytes]], bytes, float] + + +def _positive_timeout_seconds(timeout: float | httpx.Timeout) -> float: + seconds: Final = _timeout_to_seconds(timeout) + if seconds is None or not math.isfinite(seconds) or seconds <= 0: + raise ValueError("OCR bridge requires a positive finite timeout") + return seconds + + +@dataclass(kw_only=True, slots=True) +class OCRBoundary: + handler: BaseLLMHTTPHandler + model: str + document: DocumentType + optional_params: dict[str, object] + logging_obj: Logging + api_key: str | None + api_base: str | None + headers: dict[str, object] | None + provider_config: BaseOCRConfig + litellm_params: dict[str, object] + custom_llm_provider: str + timeout: float | httpx.Timeout + client: HTTPHandler | AsyncHTTPHandler | None = None + request: httpx.Request | None = field(default=None, init=False) + + def prepare(self) -> OCRRoots: + roots: Final = self.handler._prepare_ocr_request( + model=self.model, + document=self.document, + optional_params=self.optional_params, + logging_obj=self.logging_obj, + api_key=self.api_key, + api_base=self.api_base, + headers=self.headers, + provider_config=self.provider_config, + litellm_params=self.litellm_params, + ) + if not isinstance(self.client, HTTPHandler): + self.client = _get_httpx_client() + return roots + + async def aprepare(self) -> OCRRoots: + roots: Final = await self.handler._async_prepare_ocr_request( + model=self.model, + document=self.document, + optional_params=self.optional_params, + logging_obj=self.logging_obj, + api_key=self.api_key, + api_base=self.api_base, + headers=self.headers, + provider_config=self.provider_config, + litellm_params=self.litellm_params, + ) + if not isinstance(self.client, AsyncHTTPHandler): + self.client = get_async_httpx_client(llm_provider=litellm.LlmProviders(self.custom_llm_provider)) + return roots + + def encode(self, roots: OCRRoots) -> OCREncoded: + headers, url, data, _files = roots + seconds: Final = _positive_timeout_seconds(self.timeout) + if self.client is None: + raise RuntimeError("OCR boundary must be prepared before encoding") + try: + self.request = self.client.client.build_request( + "POST", + url, + headers=cast(dict[str, str], headers), + json=data, + timeout=self.timeout, + ) + return str(self.request.url), self.request.headers.raw, self.request.read(), seconds + except Exception as e: # noqa: BLE001 # match the Python OCR handler's encoding error mapping + raise self.handler._handle_error(e=e, provider_config=self.provider_config) + + def _response(self, wire: OCRWire) -> httpx.Response: + if self.request is None: + raise RuntimeError("OCR boundary must be encoded before finishing") + status, headers, content = wire + try: + response: Final = httpx.Response(status, headers=headers, content=content, request=self.request) + response.raise_for_status() + except Exception as e: # noqa: BLE001 # match the Python OCR handler's response error mapping + raise self.handler._handle_error(e=e, provider_config=self.provider_config) + return response + + def finish(self, wire: OCRWire) -> OCRResponse: + return self.handler._transform_ocr_response( + provider_config=self.provider_config, + model=self.model, + response=self._response(wire), + logging_obj=self.logging_obj, + optional_params=self.optional_params, + ) + + async def afinish(self, wire: OCRWire) -> OCRResponse: + return await self.provider_config.async_transform_ocr_response( + model=self.model, + raw_response=self._response(wire), + logging_obj=self.logging_obj, + optional_params=self.optional_params, + ) + class RustOcr(Protocol): - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise NotImplementedError + def __call__(self, boundary: OCRBoundary) -> OCRResponse: ... class RustAocr(Protocol): - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> Awaitable[dict[str, object]]: - raise NotImplementedError + def __call__(self, boundary: OCRBoundary) -> Awaitable[OCRResponse]: ... def _as_ocr(value: object) -> RustOcr | None: @@ -61,53 +162,15 @@ def load_rust_aocr() -> RustAocr | None: return _AOCR.load() -def ocr( - *, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: +def ocr(boundary: OCRBoundary) -> OCRResponse | None: rust_ocr: Final = load_rust_ocr() if rust_ocr is None: return None - return rust_ocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), - ) + return rust_ocr(boundary) -async def aocr( - *, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout: float | httpx.Timeout | None, -) -> dict[str, object] | None: +async def aocr(boundary: OCRBoundary) -> OCRResponse | None: rust_aocr: Final = load_rust_aocr() if rust_aocr is None: return None - return await rust_aocr( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - optional_params=optional_params, - timeout_seconds=_timeout_to_seconds(timeout), - ) + return await rust_aocr(boundary) diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py index 0c8b1cc2836..e38301d6779 100644 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py @@ -11,7 +11,7 @@ supplied api_base is always honoured. from litellm.llms.azure_ai.ocr.common_utils import ( is_azure_document_intelligence_model, ) -from litellm.ocr.main import _prepare_ocr_request, _rust_bridge_api_base +from litellm.ocr.main import _prepare_ocr_request _DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} _DOC_INTELLIGENCE_ENDPOINT = "https://di.cognitiveservices.azure.com" @@ -23,13 +23,6 @@ class _FakeLogging: return None -def _resolve_secret(name: str) -> str | None: - return { - "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT": _DOC_INTELLIGENCE_ENDPOINT, - "AZURE_AI_API_BASE": _AZURE_AI_API_BASE, - }.get(name) - - def _prepare(model: str, api_base: str | None): return _prepare_ocr_request( model=model, @@ -55,31 +48,59 @@ class TestIsAzureDocumentIntelligenceModel: class TestDocIntelligenceApiBaseResolution: + """Both the Python and boundary routes resolve endpoints through the provider + config's ``get_complete_url``; these pin that resolution.""" + def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): """Without an explicit api_base, the AZURE_AI_API_BASE fallback must not overwrite the endpoint, so it resolves to the Document Intelligence one.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) + monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", _DOC_INTELLIGENCE_ENDPOINT) + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) - assert prepared.api_base is None - assert _rust_bridge_api_base(prepared, _resolve_secret) == _DOC_INTELLIGENCE_ENDPOINT + url = AzureDocumentIntelligenceOCRConfig().get_complete_url( + api_base=prepared.api_base, + model=prepared.model, + optional_params=prepared.optional_params, + litellm_params=prepared.litellm_params, + ) + assert url.startswith(_DOC_INTELLIGENCE_ENDPOINT) def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): """A caller-supplied api_base must always win, even for doc-intelligence.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) custom = "https://my-di.cognitiveservices.azure.com" prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) assert prepared.api_base == custom - assert _rust_bridge_api_base(prepared, _resolve_secret) == custom + url = AzureDocumentIntelligenceOCRConfig().get_complete_url( + api_base=prepared.api_base, + model=prepared.model, + optional_params=prepared.optional_params, + litellm_params=prepared.litellm_params, + ) + assert url.startswith(custom) def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) + from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig prepared = _prepare("azure_ai/mistral-document-ai-2505", None) assert prepared.api_base == _AZURE_AI_API_BASE + url = AzureAIOCRConfig().get_complete_url( + api_base=prepared.api_base, + model=prepared.model, + optional_params=prepared.optional_params, + litellm_params=prepared.litellm_params, + ) + assert url.startswith(_AZURE_AI_API_BASE) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 36a29801e1a..9a05b248078 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,16 +1,31 @@ """Tests for the optional Rust-backed OCR path.""" +import asyncio import builtins +import contextvars +import copy +import gc import importlib +import inspect +import os +import subprocess +import sys +import threading import types -from typing import Any +import weakref +from dataclasses import replace +from datetime import datetime +from typing import Any, Final, cast import httpx import pytest import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.base_llm.ocr.transformation import OCRRequestData, OCRResponse +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.rust_bridge import configuration # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` @@ -40,101 +55,57 @@ class CapturedException(Exception): pass -class RustUpstreamError(Exception): - pass - - class RecordingBridge: - """A fake ``RustOcr`` callable that records the args it was handed.""" + """A fake ``RustOcr`` callable that records the boundary it was handed.""" def __init__(self) -> None: self.calls: list[dict[str, object]] = [] - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: + def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: self.calls.append( { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": timeout_seconds, + "model": boundary.model, + "document": boundary.document, + "api_key": boundary.api_key, + "api_base": boundary.api_base, + "custom_llm_provider": boundary.custom_llm_provider, + "extra_headers": boundary.headers, + "optional_params": boundary.optional_params, + "timeout": boundary.timeout, } ) - return dict(FAKE_OCR_RESPONSE) + return OCRResponse.model_validate(FAKE_OCR_RESPONSE) class RecordingAsyncBridge: - """A fake async ``RustAocr`` callable that records the args it was handed.""" + """A fake async ``RustAocr`` callable that records the boundary it was handed.""" def __init__(self) -> None: self.calls: list[dict[str, object]] = [] - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: + async def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: self.calls.append( { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": timeout_seconds, + "model": boundary.model, + "document": boundary.document, + "api_key": boundary.api_key, + "api_base": boundary.api_base, + "custom_llm_provider": boundary.custom_llm_provider, + "extra_headers": boundary.headers, + "optional_params": boundary.optional_params, + "timeout": boundary.timeout, } ) - return dict(FAKE_OCR_RESPONSE) + return OCRResponse.model_validate(FAKE_OCR_RESPONSE) class RaisingBridge: - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: + def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: raise RuntimeError("bridge failed") class RaisingAsyncBridge: - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - timeout_seconds: float | None, - ) -> dict[str, object]: + async def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: raise RuntimeError("bridge failed") @@ -163,6 +134,7 @@ class FakeOCRConfig: def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: self.api_key_env_var = api_key_env_var + self.seen_api_keys: list[str | None] = [] def get_api_key_env_var(self) -> str: return self.api_key_env_var @@ -176,6 +148,7 @@ class FakeOCRConfig: api_base: str | None, litellm_params: dict[str, object], ) -> dict[str, object]: + self.seen_api_keys.append(api_key) return {"Authorization": f"Bearer {api_key}", **headers} def get_complete_url( @@ -188,8 +161,36 @@ class FakeOCRConfig: ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" - def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: - return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def transform_ocr_request( + self, + *, + model: str, + document: dict[str, object], + optional_params: dict[str, object], + headers: dict[str, object], + api_key: str | None, + api_base: str | None, + ) -> OCRRequestData: + return OCRRequestData(data={"model": model, "document": document, **optional_params}, files=None) + + async def async_transform_ocr_request( + self, + *, + model: str, + document: dict[str, object], + optional_params: dict[str, object], + headers: dict[str, object], + api_key: str | None, + api_base: str | None, + ) -> OCRRequestData: + return self.transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + api_key=api_key, + api_base=api_base, + ) def build_prepared_request( @@ -206,6 +207,8 @@ def build_prepared_request( litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = 12.5, ) -> Any: + from litellm.constants import request_timeout + return ocr_main._PreparedOCRRequest( model=model, document=document, @@ -216,11 +219,24 @@ def build_prepared_request( provider_config=provider_config or FakeOCRConfig(), optional_params=optional_params or {}, litellm_params=litellm_params or {}, - effective_timeout=timeout, + effective_timeout=timeout if timeout is not None else float(request_timeout), litellm_logging_obj=logging_obj or RecordingLogging(), ) +class BoundaryDriver: + """Stands in for the native route: drives the boundary's own methods.""" + + def __init__(self) -> None: + self.roots: rust_bridge.OCRRoots | None = None + self.encoded: rust_bridge.OCREncoded | None = None + + def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: + self.roots = boundary.prepare() + self.encoded = boundary.encode(self.roots) + return OCRResponse.model_validate(FAKE_OCR_RESPONSE) + + @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" @@ -395,86 +411,20 @@ def test_timeout_to_seconds_handles_float_timeout_and_none(): assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): - bridge = RecordingBridge() - - litellm.rust(True) - - rust_bridge._OCR.override(bridge) - response = rust_bridge.ocr( - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout=12.5, - ) - - assert response == FAKE_OCR_RESPONSE - call = bridge.calls[0] - assert call == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True, "pages": [0]}, - "timeout_seconds": 12.5, - } - - -@pytest.mark.asyncio -async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): - bridge = RecordingAsyncBridge() - - litellm.rust(True) - - rust_bridge._AOCR.override(bridge) - response = await rust_bridge.aocr( - model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout=httpx.Timeout(30.0, read=42.0), - ) - - assert response == FAKE_OCR_RESPONSE - assert bridge.calls[0] == { - "model": "mistral-ocr-maas", - "document": DOCUMENT, - "api_key": None, - "api_base": None, - "custom_llm_provider": "vertex_ai", - "extra_headers": None, - "optional_params": {"vertex_project": "project-1"}, - "timeout_seconds": 42.0, - } - - -def test_run_rust_ocr_prepares_request_and_wraps_response(): +def test_run_rust_ocr_forwards_boundary_fields(): bridge = RecordingBridge() logging_obj = RecordingLogging() litellm.rust(True) rust_bridge._OCR.override(bridge) response = ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", extra_headers={"x-trace-id": "trace-1"}, optional_params={"include_image_base64": True}, timeout=12.5, - ), - resolve_api_key=lambda _name: None, + ) ) assert isinstance(response, OCRResponse) @@ -485,202 +435,76 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): "api_key": "sk-test", "api_base": "https://proxy.internal", "custom_llm_provider": "mistral", - "extra_headers": { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - }, + "extra_headers": {"x-trace-id": "trace-1"}, "optional_params": {"include_image_base64": True}, - "timeout_seconds": 12.5, + "timeout": 12.5, } + assert logging_obj.pre_call_kwargs is None # pre_call now runs inside the boundary -def test_rust_upstream_error_uses_ocr_provider_error_mapping(): - error = RustUpstreamError(400, '{"message":"invalid model"}') - - mapped = ocr_main._map_rust_ocr_error( - error, - build_prepared_request(), - (RuntimeError, RustUpstreamError), - ) - - assert isinstance(mapped, BaseLLMException) - assert mapped.status_code == 400 - assert mapped.message == '{"message":"invalid model"}' - - -def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - bridge = RecordingBridge() +def test_run_rust_ocr_passes_raw_api_key_to_provider_config(): + """Key resolution moved into provider ``validate_environment``; the boundary + forwards the caller's key unchanged.""" + provider_config = FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY") + driver = BoundaryDriver() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.override(driver) - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request(api_key=None, timeout=None), - resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, - ) + ocr_main._run_rust_ocr(build_prepared_request(provider_config=provider_config, api_key="sk-explicit", timeout=None)) + ocr_main._run_rust_ocr(build_prepared_request(provider_config=provider_config, api_key=None, timeout=None)) - assert bridge.calls[0]["api_key"] == "sk-from-vault" + assert provider_config.seen_api_keys == ["sk-explicit", None] -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() +def test_run_rust_ocr_preserves_native_response_identity(): + """The bridge returns the boundary's own finish() object, not a re-validated copy.""" + sentinel = OCRResponse.model_validate(FAKE_OCR_RESPONSE) + + class IdentityBridge: + def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: + return sentinel + litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.override(IdentityBridge()) - def _resolver(name: str) -> str | None: - raise AssertionError(f"resolver should not be called for {name}") + response = ocr_main._run_rust_ocr(build_prepared_request()) - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - api_key="sk-explicit", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" + assert response is sentinel -def test_run_rust_ocr_uses_provider_api_key_env_var(): - bridge = RecordingBridge() - resolver_calls = [] - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name): - resolver_calls.append(name) - return "sk-provider-env" - - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), - model="provider-ocr-model", - api_key=None, - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert resolver_calls == ["PROVIDER_OCR_API_KEY"] - assert bridge.calls[0]["api_key"] == "sk-provider-env" - - -def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - litellm_params={ - "vertex_project": "project-1", - "vertex_location": "us-central1", - "vertex_credentials": "redacted", - }, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["optional_params"] == { - "include_image_base64": True, - "vertex_project": "project-1", - "vertex_location": "us-central1", - } - - -def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name: str) -> str | None: - return { - "VERTEXAI_PROJECT": "project-from-secret", - "VERTEXAI_LOCATION": "us-east5", - }.get(name) - - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" - assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" - - -def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, - ) - - assert bridge.calls[0]["api_base"] == "https://azure.example.com" - - -def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - custom_llm_provider="azure_ai", - model="doc-intelligence/prebuilt-layout", - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None - ), - ) - - assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" - - -def test_run_rust_ocr_runs_pre_call_logging(): +def test_boundary_prepare_runs_pre_call_and_encodes_the_same_roots(): + """The logging view must alias the execution roots: mutating the headers the + callback received must surface in the encoded wire headers, while replacing + a view field must not.""" logging_obj = RecordingLogging() - bridge = RecordingBridge() + driver = BoundaryDriver() litellm.rust(True) - rust_bridge._OCR.override(bridge) + rust_bridge._OCR.override(driver) - ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - logging_obj=logging_obj, - api_base="https://api.mistral.ai/v1", - extra_headers={"x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) + seen: dict[str, object] = {} + + original_pre_call = logging_obj.pre_call + + def observing_pre_call(**kwargs: object) -> None: + original_pre_call(**kwargs) + view = cast(dict[str, object], kwargs["additional_args"]) + seen["headers"] = view["headers"] + cast(dict[str, object], view["headers"])["X-Proof"] = "mutated" + + logging_obj.pre_call = observing_pre_call # type: ignore[method-assign] + + ocr_main._run_rust_ocr(build_prepared_request(logging_obj=logging_obj, api_base="https://api.mistral.ai/v1")) assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" - assert additional_args["headers"] == { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - } + assert additional_args["headers"] == {"Authorization": "Bearer sk-test", "X-Proof": "mutated"} + assert driver.roots is not None + roots_headers, _url, _data, _files = driver.roots + assert roots_headers is seen["headers"] + assert driver.encoded is not None + encoded_headers = {name.lower(): value for name, value in driver.encoded[1]} + assert encoded_headers[b"x-proof"] == b"mutated" def test_ocr_routes_to_rust_when_enabled(fake_bridge): @@ -700,10 +524,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" assert call["custom_llm_provider"] == "mistral" - assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - } + assert call["extra_headers"] == {"x-trace-id": "trace-1"} assert call["optional_params"].get("include_image_base64") is True @@ -772,10 +593,7 @@ async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" assert call["custom_llm_provider"] == "mistral" - assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - } + assert call["extra_headers"] == {"x-trace-id": "trace-1"} assert call["optional_params"].get("include_image_base64") is True @@ -805,7 +623,7 @@ def test_ocr_forwards_timeout_to_rust(fake_bridge): client ceiling doesn't silently override shorter deadlines.""" litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) - assert fake_bridge.calls[0]["timeout_seconds"] == 12.5 + assert fake_bridge.calls[0]["timeout"] == 12.5 def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): @@ -813,7 +631,7 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): from litellm.constants import request_timeout - assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) + assert fake_bridge.calls[0]["timeout"] == float(request_timeout) def test_ocr_does_not_route_to_rust_when_disabled(): @@ -866,3 +684,521 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" + + +################################################# +# Proof: pre_call callbacks run against the same objects the wire request is +# built from, through the real native route, compared with the Python route. +################################################# + +MISTRAL_OCR_RESPONSE_JSON: Final = ( + b'{"pages": [{"index": 0, "markdown": "proof"}], "model": "mistral-ocr-2505-completion",' + b' "document_annotation": null, "usage_info": {"pages_processed": 1}, "object": "ocr"}' +) + + +class _WireRecorder: + """Loopback OCR server recording every request it serves.""" + + def __init__(self) -> None: + import json + import threading + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + self.requests: list[dict[str, object]] = [] + self.received = threading.Event() + self.release = threading.Event() + self.finished = threading.Event() + self.release.set() + self.status = 200 + recorder = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # http.server API + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + recorder.requests.append( + { + "path": self.path, + "headers": {name.lower(): value for name, value in self.headers.items()}, + "body": json.loads(body), + } + ) + recorder.received.set() + try: + if not recorder.release.wait(timeout=10): + return + self.send_response(recorder.status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(MISTRAL_OCR_RESPONSE_JSON))) + self.end_headers() + self.wfile.write(MISTRAL_OCR_RESPONSE_JSON) + except (BrokenPipeError, ConnectionResetError): + pass + finally: + recorder.finished.set() + + def log_message(self, *args: object) -> None: + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + + @property + def api_base(self) -> str: + host, port = self._server.server_address[:2] + return f"http://{host}:{port}" + + def stop(self) -> None: + self.release.set() + self._server.shutdown() + self._server.server_close() + self._thread.join(timeout=5) + + +@pytest.fixture +def wire_recorder(): + recorder = _WireRecorder() + try: + yield recorder + finally: + recorder.stop() + + +def _native_boundary_route_available() -> bool: + bridge = rust_bridge_loader.get_native_bridge() + if bridge is None: + return False + ocr_fn = getattr(bridge, "ocr", None) + if ocr_fn is None: + return False + try: + return str(inspect.signature(ocr_fn)) == "(boundary)" + except (TypeError, ValueError): + return False + + +@pytest.fixture +def native_ocr(): + if not _native_boundary_route_available(): + if os.environ.get("LITELLM_REQUIRE_NATIVE_OCR") == "1": + pytest.fail("native OCR boundary route not built") + pytest.skip("native OCR boundary route not built") + return rust_bridge_loader.get_native_bridge() + + +class ProofCallback(CustomLogger): + def __init__(self, document, context, events, fail=False) -> None: + self.document = document + self.context = context + self.events = events + self.fail = fail + self.headers: dict[str, object] | None = None + self.body: dict[str, object] | None = None + self.details = None + self.calls = 0 + + def log_pre_api_call(self, model, messages, kwargs): + self.calls += 1 + self.events.append(("mutate", self.context.get(), threading.get_ident(), asyncio.current_task())) + self.context.set("callback") + self.details = kwargs + view = kwargs["additional_args"] + self.headers = view["headers"] + self.body = view["complete_input_dict"] + self.headers["X-Proof"] = "mutated" + self.body["document"]["document_url"] = "https://example.invalid/mutated-by-callback.pdf" + self.document["document_name"] = "closure-mutation" + view["headers"] = {"X-Replacement": "must-not-reach-wire"} + view["complete_input_dict"] = {"model": "logging-only"} + if self.fail: + raise RuntimeError("expected pre-call failure") + return {"additional_args": {"headers": {"X-Return": "ignored"}}} + + +def _assert_retained_wire(request: dict[str, object]) -> None: + assert request["path"] == "/v1/ocr" + headers = request["headers"] + assert headers["authorization"] == "Bearer sk-test" + assert headers["x-proof"] == "mutated" + assert "x-replacement" not in headers + assert "x-return" not in headers + body = request["body"] + assert body["model"] == "mistral-ocr-latest" + assert body["document"]["document_url"] == "https://example.invalid/mutated-by-callback.pdf" + assert body["document"]["document_name"] == "closure-mutation" + + +async def _pre_call_contract(native, wire_recorder, monkeypatch, asynchronous, enabled, fail=False, control="original"): + document = { + "type": "document_url", + "document_url": "https://example.invalid/original.pdf", + } + context = contextvars.ContextVar("ocr-pre-call", default="caller") + events = [] + observations = [] + native_calls = [] + proof = ProofCallback(document, context, events, fail=fail) + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + events.append(("observe", context.get(), threading.get_ident(), asyncio.current_task())) + view = kwargs["additional_args"] + observations.append( + ( + kwargs is proof.details, + tuple(view["headers"].items()), + view["complete_input_dict"]["model"], + document["document_url"], + ) + ) + + def selected(boundary): + if control == "copy-input": + return replace(boundary, document=copy.deepcopy(boundary.document)) + if control in {"logging-roots", "tuple-only"}: + + class EncodeControl: + def __getattr__(self, name): + return getattr(boundary, name) + + def encode(self, roots): + headers, url, body, files = roots + if control == "logging-roots": + view = boundary.logging_obj.model_call_details["additional_args"] + return boundary.encode((view["headers"], url, view["complete_input_dict"], files)) + return boundary.encode((headers, url, body, files)) + + return EncodeControl() + return boundary + + def sync_call(boundary): + native_calls.append("sync") + if control == "duplicate-prepare": + boundary.prepare() + return native.ocr(selected(boundary)) + + async def async_call(boundary): + native_calls.append("async") + if control == "duplicate-prepare": + await boundary.aprepare() + if control == "new-task": + return await asyncio.create_task(native.aocr(selected(boundary))) + return await native.aocr(selected(boundary)) + + rust_bridge._OCR.override(sync_call) + rust_bridge._AOCR.override(async_call) + monkeypatch.setattr(litellm, "input_callback", [proof, Observe()]) + litellm.rust(enabled) + caller = (threading.get_ident(), asyncio.current_task()) + arguments = dict(model=MODEL, document=document, api_key="sk-test", api_base=wire_recorder.api_base, num_retries=0) + response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert native_calls == (["async" if asynchronous else "sync"] if enabled else []), "native dispatch" + assert isinstance(response, OCRResponse) and response.pages[0].markdown == "proof" + assert proof.calls == 1, "callback count" + assert proof.body is not None and proof.body["document"] is document, "caller identity" + assert events == [("mutate", "caller", *caller), ("observe", "callback", *caller)], "callback context/order" + assert context.get() == "callback", "caller context write" + assert observations == [ + ( + True, + (("X-Replacement", "must-not-reach-wire"),), + "logging-only", + "https://example.invalid/mutated-by-callback.pdf", + ) + ], "observation-time values" + assert len(wire_recorder.requests) == 1, "POST count" + assert wire_recorder.requests[0]["headers"].get("x-proof") == "mutated", "execution roots" + _assert_retained_wire(wire_recorder.requests[0]) + proof.body["document"]["document_url"] = "https://example.invalid/after-encode.pdf" + assert document["document_url"] == "https://example.invalid/after-encode.pdf" + assert ( + wire_recorder.requests[0]["body"]["document"]["document_url"] + == "https://example.invalid/mutated-by-callback.pdf" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +@pytest.mark.parametrize("fail", [False, True], ids=["return-ignored", "caught-error"]) +async def test_pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled, fail): + await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled, fail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "asynchronous, control, message", + [ + (False, "copy-input", "caller identity"), + (True, "copy-input", "caller identity"), + (False, "duplicate-prepare", "callback count"), + (True, "duplicate-prepare", "callback count"), + (True, "new-task", "callback context/order"), + (False, "logging-roots", "execution roots"), + (True, "logging-roots", "execution roots"), + ], +) +async def test_pre_call_contract_rejects_boundary_mutants( + native_ocr, wire_recorder, monkeypatch, asynchronous, control, message +): + with pytest.raises(AssertionError, match=message): + await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, True, control=control) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_pre_call_contract_allows_root_tuple_reconstruction(native_ocr, wire_recorder, monkeypatch, asynchronous): + await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, True, control="tuple-only") + + +class PreCallAbort(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +async def test_pre_call_escape_never_sends_or_replays(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled): + document = {"type": "document_url", "document_url": "https://example.invalid/original.pdf"} + error = PreCallAbort("stop before POST") + calls = [] + native_calls = [] + + class Abort(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + calls.append(kwargs["additional_args"]["complete_input_dict"]["document"]) + document["document_url"] = "https://example.invalid/aborted.pdf" + raise error + + def sync_call(boundary): + native_calls.append("sync") + return native_ocr.ocr(boundary) + + async def async_call(boundary): + native_calls.append("async") + return await native_ocr.aocr(boundary) + + rust_bridge._OCR.override(sync_call) + rust_bridge._AOCR.override(async_call) + monkeypatch.setattr(litellm, "input_callback", [Abort()]) + litellm.rust(enabled) + arguments = dict(model=MODEL, document=document, api_key="sk-test", api_base=wire_recorder.api_base, num_retries=0) + with pytest.raises(PreCallAbort, match="stop before POST") as caught: + await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert caught.value is error + assert len(calls) == 1 and calls[0] is document + assert native_calls == (["async" if asynchronous else "sync"] if enabled else []) + assert document["document_url"] == "https://example.invalid/aborted.pdf" + assert wire_recorder.requests == [] + + +class RetainedDocument(dict): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +@pytest.mark.parametrize("outcome", ["success", "error", "cancel"]) +async def test_ocr_retention_during_post_and_terminal_cleanup(native_ocr, wire_recorder, enabled, outcome): + retained = [] + references = [] + calls = [] + wire_recorder.release.clear() + wire_recorder.status = 429 if outcome == "error" else 200 + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + body = kwargs["additional_args"]["complete_input_dict"] + retained.append(body) + references.append(weakref.ref(body["document"])) + calls.append("pre_call") + + async def request(): + document = RetainedDocument(type="document_url", document_url="https://example.invalid/original.pdf") + logging_obj = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id="retention", + function_id="retention", + dynamic_input_callbacks=[Retain()], + ) + arguments = dict( + model="mistral-ocr-latest", + document=document, + optional_params={}, + logging_obj=logging_obj, + api_key="sk-test", + api_base=wire_recorder.api_base, + headers=None, + provider_config=MistralOCRConfig(), + litellm_params={}, + custom_llm_provider="mistral", + timeout=5.0, + ) + if enabled: + return await native_ocr.aocr(rust_bridge.OCRBoundary(handler=ocr_main.base_llm_http_handler, **arguments)) + return await ocr_main.base_llm_http_handler.async_ocr(**arguments) + + task = asyncio.create_task(request()) + try: + assert await asyncio.to_thread(wire_recorder.received.wait, 5), "POST never reached server" + assert not task.done() + assert calls == ["pre_call"] + assert len(references) == 1 and references[0]() is not None + retained[0]["document"]["document_url"] = "https://example.invalid/after-consumption.pdf" + retained[0]["document"]["cycle"] = retained[0] + assert wire_recorder.requests[0]["body"]["document"]["document_url"] == "https://example.invalid/original.pdf" + if outcome == "cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + elif outcome == "error": + wire_recorder.release.set() + with pytest.raises(BaseLLMException) as caught: + await task + assert caught.value.status_code == 429 + del caught + else: + wire_recorder.release.set() + response = await task + assert response.pages[0].markdown == "proof" + finally: + wire_recorder.release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert await asyncio.to_thread(wire_recorder.finished.wait, 5), "server did not finish" + del task + gc.collect() + assert references[0]() is retained[0]["document"] + assert retained[0]["document"]["document_url"] == "https://example.invalid/after-consumption.pdf" + assert len(wire_recorder.requests) == 1 and calls == ["pre_call"] + retained.clear() + await asyncio.sleep(0) + gc.collect() + assert references[0]() is None, "execution retained the caller graph after cleanup" + + +@pytest.mark.parametrize("filename", ["ocr_driver.py", "retained_callback.py"]) +def test_native_ocr_cold_cache_reentry(native_ocr, filename): + script = """ +import sys +from litellm.rust_bridge import _native + +events = [] +compilations = [] +error = LookupError('preparation stopped') + +class Boundary: + async def aprepare(self): + raise error + +def invoke(): + pending = _native.aocr(Boundary()) + try: + pending.send(None) + except LookupError as caught: + assert caught is error + events.append('raised') + else: + raise AssertionError('preparation did not raise') + finally: + pending.close() + error.__traceback__ = None + +def audit(event, args): + if event == 'compile' and args[1] == sys.argv[1]: + compilations.append(args[1]) + if len(compilations) == 1: + events.append('entered') + invoke() + events.append('returned') + +sys.addaudithook(audit) +invoke() +invoke() +assert events == ['entered', 'raised', 'returned', 'raised', 'raised'], events +assert len(compilations) == 2, compilations +print('cold reentry passed') +""" + isolation = ["-I"] if sys.flags.isolated else [] + result = subprocess.run( + [sys.executable, *isolation, "-c", script, filename], capture_output=True, text=True, timeout=30 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "cold reentry passed" + + +@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) +def test_sync_pre_call_reentry_without_event_loop(native_ocr, wire_recorder, monkeypatch, enabled): + context = contextvars.ContextVar("sync-ocr-context", default="caller") + events = [] + native_calls = [] + document = {"type": "document_url", "document_url": "https://example.invalid/original.pdf"} + + class Reenter(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + try: + asyncio.get_running_loop() + except RuntimeError: + loop_running = False + else: + loop_running = True + events.append((context.get(), threading.get_ident(), loop_running)) + if len(events) == 1: + context.set("nested") + response = litellm.ocr( + model=MODEL, + document=dict(document), + api_key="sk-test", + api_base=wire_recorder.api_base, + num_retries=0, + ) + events.append((response.pages[0].markdown, threading.get_ident(), loop_running)) + + def sync_call(boundary): + native_calls.append(boundary) + return native_ocr.ocr(boundary) + + rust_bridge._OCR.override(sync_call) + monkeypatch.setattr(litellm, "input_callback", [Reenter()]) + litellm.rust(enabled) + response = litellm.ocr( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + num_retries=0, + ) + assert response.pages[0].markdown == "proof" + assert events == [(value, threading.get_ident(), False) for value in ("caller", "nested", "proof")] + assert context.get() == "nested" + assert len(native_calls) == (2 if enabled else 0) + assert len(wire_recorder.requests) == 2 + + +@pytest.mark.parametrize("explicit_close", [False, True], ids=["abandoned", "closed"]) +def test_unstarted_native_ocr_driver_releases_cyclic_input(native_ocr, explicit_close): + document = RetainedDocument(type="document_url", document_url="https://example.invalid/original.pdf") + reference = weakref.ref(document) + logger = RecordingLogging() + boundary = ocr_main._ocr_boundary(build_prepared_request(document=document, logging_obj=logger)) + pending = native_ocr.aocr(boundary) + document["pending"] = pending + del boundary, document + assert reference() is not None + assert logger.pre_call_kwargs is None + if explicit_close: + pending.close() + del pending + gc.collect() + else: + del pending + with pytest.warns(RuntimeWarning, match="coroutine .* was never awaited"): + gc.collect() + assert reference() is None + assert logger.pre_call_kwargs is None diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index a7f50a82a99..4f5aed50985 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -10,11 +10,13 @@ import sys import tempfile import threading import zipfile +from dataclasses import dataclass from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from socket import socket as Socket from typing import Final +from urllib.error import HTTPError REQUEST_STARTED: Final = threading.Event() REQUEST_CANCELLED: Final = threading.Event() @@ -123,20 +125,54 @@ def load_native(native_path: Path) -> object: return native_module +@dataclass(frozen=True) +class OCRBoundary: + api_base: str + outcome: str + + def prepare(self) -> dict[str, object]: + return { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "include_image_base64": True, + } + + async def aprepare(self) -> dict[str, object]: + return self.prepare() + + def encode(self, roots: dict[str, object]) -> tuple[str, list[tuple[bytes, bytes]], bytes, float]: + return ( + f"{self.api_base}/v1/ocr", + [ + (b"authorization", b"Bearer sk-native"), + (b"content-type", b"application/json"), + (b"x-test-route", b"ocr"), + (b"x-test-outcome", self.outcome.encode()), + ], + json.dumps(roots).encode(), + 3.0, + ) + + def finish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object: + status, headers, content = wire + assert (b"content-type", b"application/json") in headers + if status != 200: + assert content == b'{"error":"native-rate-limit"}' + raise HTTPError(f"{self.api_base}/v1/ocr", status, "native-rate-limit", HTTPMessage(), None) + return json.loads(content) + + async def afinish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object: + return self.finish(wire) + + def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: + if route == "ocr": + return {"boundary": OCRBoundary(api_base, outcome)} common: Final = { "api_base": api_base, "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, "timeout_seconds": 3.0, } - if route == "ocr": - return common | { - "model": "mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, - "api_key": "sk-native", - "custom_llm_provider": "mistral", - "optional_params": {"include_image_base64": True}, - } if route == "transcription": return common | { "model": "mistral.voxtral-mini-3b-2507", @@ -192,7 +228,11 @@ def success_value(route: str, response: dict[object, object]) -> object: def assert_rate_limit(native: object, route: str, error: BaseException) -> None: - if route in {"ocr", "chat_completions"}: + if route == "ocr": + if not isinstance(error, HTTPError) or error.code != 429: + raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + return + if route == "chat_completions": upstream_error: Final = native.RustUpstreamError if not isinstance(error, upstream_error) or error.args[0] != 429: raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") @@ -207,7 +247,7 @@ def exercise_sync(native: object, api_base: str) -> None: assert_success(route, function(**route_kwargs(route, api_base, "success"))) try: function(**route_kwargs(route, api_base, "429")) - except (RuntimeError, native.RustUpstreamError) as error: + except (HTTPError, RuntimeError, native.RustUpstreamError) as error: assert_rate_limit(native, route, error) else: raise AssertionError(f"{route} accepted a 429 response") @@ -219,7 +259,7 @@ async def exercise_async(native: object, api_base: str) -> None: assert_success(route, await function(**route_kwargs(route, api_base, "success"))) try: await function(**route_kwargs(route, api_base, "429")) - except (RuntimeError, native.RustUpstreamError) as error: + except (HTTPError, RuntimeError, native.RustUpstreamError) as error: assert_rate_limit(native, route, error) else: raise AssertionError(f"a{route} accepted a 429 response") @@ -227,12 +267,7 @@ async def exercise_async(native: object, api_base: str) -> None: async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather( - *( - native.amessages(**route_kwargs("messages", api_base, "success")) - for _ in range(32) - ) - ), + asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: