diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 45558be9a8c..aaf31514679 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,7 +4,6 @@ on: push: paths: - "litellm-rust/**" - - "tests/test_litellm/**" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -12,6 +11,7 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" + - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,7 +21,6 @@ on: - "litellm_**" paths: - "litellm-rust/**" - - "tests/test_litellm/**" - ".cargo/**" - "pyproject.toml" - "uv.lock" @@ -29,6 +28,7 @@ on: - ".github/actions/setup-uv-with-retries/**" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" + - "tests/test_litellm/rust_bridge/native_route_wheel_test.py" - ".github/workflows/test-rust.yml" permissions: @@ -102,7 +102,7 @@ jobs: restore-keys: | ${{ runner.os }}-cargo-${{ github.job }}- - - run: cargo test --workspace --locked --exclude litellm-python-interop --exclude litellm-python-bridge + - run: cargo test --workspace --locked working-directory: litellm-rust - run: cargo test -p litellm-core --features bedrock-auth --locked @@ -131,5 +131,5 @@ jobs: - name: Check Python fixtures for Cargo tests run: make lint-rust-python-fixtures - - name: Test retained callbacks and OCR with repo Python + - name: Test retained callbacks with repo Python run: make test-rust-python diff --git a/Makefile b/Makefile index c70cb026e8f..4fd2b23d53f 100644 --- a/Makefile +++ b/Makefile @@ -314,11 +314,11 @@ test-rust-python: install-rust-python-test-deps PYTHONPATH="$(CURDIR):$$site_packages$${PYTHONPATH:+:$$PYTHONPATH}" \ LITELLM_LOCAL_MODEL_COST_MAP=True \ cargo test --manifest-path litellm-rust/Cargo.toml \ - -p litellm-python-interop -p litellm-python-bridge --tests --locked -- --include-ignored + -p litellm-python-interop --tests --locked -- --include-ignored lint-rust-python-fixtures: - $(UV) tool run --from ruff==0.15.3 ruff check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests - $(UV) tool run --from ruff==0.15.3 ruff format --check --config ruff-tests.toml litellm-rust/crates/python-interop/tests litellm-rust/crates/python-bridge/tests + $(UV) tool run --from ruff==0.15.3 ruff check --config ruff-tests.toml litellm-rust/crates/python-interop/tests + $(UV) tool run --from ruff==0.15.3 ruff format --check --config ruff-tests.toml litellm-rust/crates/python-interop/tests test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 54757ae634f..c2c1e454e6a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1471,10 +1471,8 @@ dependencies = [ "litellm-python-interop", "pyo3", "pyo3-async-runtimes", - "rstest", "serde", "serde_json", - "serial_test", "tokio", "tokio-tungstenite", "tracing", diff --git a/litellm-rust/README.md b/litellm-rust/README.md index a5629744574..74f28f61680 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -58,8 +58,8 @@ runs for changes under `litellm-rust/`. ### Python-Integrated Tests From the repository root, run the ignored Cargo tests that need the repository's -Python dependencies and the pinned Ruff checks over both crates' `tests/` -directories: +Python dependencies and the pinned Ruff checks over the interop crate's Python +test fixtures: ```bash make test-rust-python @@ -68,10 +68,16 @@ make lint-rust-python-fixtures `test-rust-python` installs the locked SDK dependencies with uv, points `PYO3_PYTHON` at the project interpreter, and runs -`cargo test -p litellm-python-interop -p litellm-python-bridge --tests --locked`. -`lint-rust-python-fixtures` runs pinless linting only, no sync. +`cargo test -p litellm-python-interop --tests --locked -- --include-ignored`. +`lint-rust-python-fixtures` runs pinned Ruff lint and formatting checks without +syncing the project environment -The callback lifecycle and retained OCR scenarios use +These 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 + +The callback lifecycle scenarios use `#[serial(python_interpreter)]` to isolate CPython GC and interpreter-wide LiteLLM settings under `cargo test`. Compatible tests in the same binary use `#[parallel(python_interpreter)]`: they may overlap each other, but not an diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index ec1a2988e75..fc81f4fa029 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -32,8 +32,6 @@ 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/mod.rs b/litellm-rust/crates/core/src/http_utils.rs similarity index 99% rename from litellm-rust/crates/core/src/http_utils/mod.rs rename to litellm-rust/crates/core/src/http_utils.rs index 93dfc1a5a8f..cb472dd5a57 100644 --- a/litellm-rust/crates/core/src/http_utils/mod.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -1,7 +1,5 @@ //! 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 deleted file mode 100644 index 3272beaa7e6..00000000000 --- a/litellm-rust/crates/core/src/http_utils/buffered_post.rs +++ /dev/null @@ -1,77 +0,0 @@ -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, - }) -} - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs b/litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs deleted file mode 100644 index 5237e91b66e..00000000000 --- a/litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs +++ /dev/null @@ -1,44 +0,0 @@ -use super::*; - -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/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index ccf8e8bd7a3..42282ca4da4 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,6 +1,3 @@ litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop. Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. - -The retained routes (`ocr_retained`, and any future `retained_http` variants) are a Python-compatibility adapter. Python owns prepare, transform, and logging; Rust owns only the buffered POST and the boundary call marshaling. There is no retry, billing, guardrail, or logging orchestration here, and none of it should be added. - diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 593f6d3cfca..bda09a7d840 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [lib] name = "_native" -crate-type = ["cdylib", "rlib"] +crate-type = ["cdylib"] [features] default = ["abi3"] @@ -34,8 +34,6 @@ tokio.workspace = true [dev-dependencies] criterion = "0.8.2" -rstest.workspace = true -serial_test.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index c3da372ed14..f3648158cf6 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -37,37 +37,6 @@ fn run_sync_on( where T: Serialize + Send + 'static, F: Future> + Send + 'static, -{ - let result = run_sync_value_on(py, runtime, 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, - map_error, - ) -} - -fn run_sync_value_on( - py: Python<'_>, - runtime: &Runtime, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult -where - T: Send + 'static, - F: Future> + Send + 'static, { if Handle::try_current().is_ok() { return Err(PyRuntimeError::new_err( @@ -76,7 +45,8 @@ where } let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?; - map_core_result(result, map_error) + let result = map_core_result(result, map_error)?; + Pythonized(result).into_pyobject(py).map(Bound::unbind) } pub(crate) fn run_async( @@ -89,19 +59,12 @@ where F: Future> + Send + 'static, { pyo3_async_runtimes::tokio::future_into_py(py, async move { - let result = run_async_value(future, map_error).await?; + let result = catch_future_panic(future).await?; + let result = map_core_result(result, map_error)?; Ok(Pythonized(result)) }) } -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), @@ -112,7 +75,7 @@ fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) - } } -pub(crate) async fn catch_future_panic(future: F) -> PyResult> +async fn catch_future_panic(future: F) -> PyResult> where F: Future>, { diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c00319eb83e..384f0be5a1b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -63,7 +63,7 @@ impl ResponsesWebSocketConnection { } #[pymodule(gil_used = false)] -pub mod _native { +mod _native { use pyo3::prelude::*; #[pymodule_init] @@ -98,8 +98,6 @@ mod tests { "RustUpstreamError", "ocr", "aocr", - "ocr_retained", - "aocr_retained", "transcription", "atranscription", "messages", diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 2873a9a8302..7e81f2ffe9b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -10,12 +10,9 @@ mod audio_transcription; mod chat_completions; mod messages; mod ocr; -mod ocr_retained; -mod retained_http; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { ocr::register(module)?; - ocr_retained::register(module)?; audio_transcription::register(module)?; messages::register(module)?; chat_completions::register(module)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs b/litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs deleted file mode 100644 index 9b36878013c..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs +++ /dev/null @@ -1,18 +0,0 @@ -use pyo3::prelude::*; - -use super::retained_http; - -#[pyfunction] -fn ocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult> { - retained_http::run_sync(boundary) -} - -#[pyfunction] -fn aocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult> { - retained_http::run_async(boundary) -} - -pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(ocr_retained, module)?)?; - module.add_function(wrap_pyfunction!(aocr_retained, module)?) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/retained_http.rs b/litellm-rust/crates/python-bridge/src/routes/retained_http.rs deleted file mode 100644 index 53d4d6ced42..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/retained_http.rs +++ /dev/null @@ -1,203 +0,0 @@ -use litellm_core::http_utils::buffered_post::{self, Request, Response}; -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; -use pyo3::prelude::*; -use pyo3::sync::PyOnceLock; -use pyo3::types::{PyBytes, PyList, PyTuple}; - -use crate::errors::core_error_to_pyerr; -use crate::execution::{run_async_value, run_sync_value}; - -pub(crate) struct MethodBinding { - pub(crate) name: &'static str, - pub(crate) mode: InvocationMode, -} - -pub(crate) enum BoundaryMethod { - Prepare, - Encode, - Finish, -} - -impl BoundaryMethod { - pub(crate) fn resolve(self, asynchronous: bool) -> MethodBinding { - match (self, asynchronous) { - (Self::Prepare, true) => MethodBinding { - name: "aprepare", - mode: InvocationMode::Await, - }, - (Self::Prepare, false) => MethodBinding { - name: "prepare", - mode: InvocationMode::Direct, - }, - (Self::Encode, _) => MethodBinding { - name: "encode", - mode: InvocationMode::Direct, - }, - (Self::Finish, true) => MethodBinding { - name: "afinish", - mode: InvocationMode::Await, - }, - (Self::Finish, false) => MethodBinding { - name: "finish", - mode: InvocationMode::Direct, - }, - } - } -} - -fn invoke( - boundary: &Bound<'_, PyAny>, - binding: MethodBinding, - args: Bound<'_, PyTuple>, -) -> PyResult> { - let call = PreparedCall::new( - binding.mode, - boundary.getattr(binding.name)?.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> { - invoke( - boundary, - BoundaryMethod::Prepare.resolve(asynchronous), - PyTuple::empty(boundary.py()), - ) -} - -#[pyfunction] -fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult> { - invoke( - boundary, - BoundaryMethod::Encode.resolve(false), - PyTuple::new(boundary.py(), [roots])?, - ) -} - -fn request(execution: &Bound<'_, PyAny>) -> PyResult { - type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; - let encoded = execution.call_method0("encode")?; - let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) = - encoded.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, - }) -} - -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<'py>(execution: &Bound<'py, PyAny>) -> PyResult> { - let request = request(execution)?; - pyo3_async_runtimes::tokio::future_into_py(execution.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> { - invoke( - boundary, - BoundaryMethod::Finish.resolve(asynchronous), - PyTuple::new(boundary.py(), [wire])?, - ) -} - -pub(crate) fn run_sync(boundary: &Bound<'_, PyAny>) -> PyResult> { - let py = boundary.py(); - let execution = execution_module(py)? - .getattr("RetainedExecution")? - .call1((boundary,))?; - execution.call_method0("prepare")?; - let request = request(&execution)?; - let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?; - let wire = Wire(response).into_pyobject(py)?; - execution.call_method1("finish", (wire,)).map(Bound::unbind) -} - -pub(crate) fn run_async(boundary: &Bound<'_, PyAny>) -> PyResult> { - let module = execution_module(boundary.py())?; - let execution = module.getattr("RetainedExecution")?.call1((boundary,))?; - module - .getattr("drive")? - .call1((execution,)) - .map(Bound::unbind) -} - -fn execution_module(py: Python<'_>) -> PyResult> { - static MODULE: PyOnceLock> = PyOnceLock::new(); - if MODULE.get(py).is_none() { - let module = PyModule::from_code( - py, - c"class RetainedExecution: - __slots__ = ('binding', 'roots') - - def __init__(self, binding): - self.binding = binding - self.roots = None - - def prepare(self): - self.roots = _prepare(self.binding, False) - - async def aprepare(self): - self.roots = await _prepare(self.binding, True) - - def encode(self): - return _encode(self.binding, self.roots) - - def finish(self, wire): - return _finish(self.binding, wire, False) - - async def afinish(self, wire): - return await _finish(self.binding, wire, True) - -async def drive(execution): - await execution.aprepare() - wire = await _send(execution) - return await execution.afinish(wire) -", - c"retained_execution.py", - c"_retained_execution", - )?; - module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; - module.add("_encode", wrap_pyfunction!(encode, &module)?)?; - module.add("_finish", wrap_pyfunction!(finish, &module)?)?; - module.add("_send", wrap_pyfunction!(send, &module)?)?; - let _ = MODULE.set(py, module.unbind()); - } - let module = MODULE.get(py).unwrap(); - Ok(module.bind(py).clone()) -} diff --git a/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py b/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py deleted file mode 100644 index 94daa1ce969..00000000000 --- a/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py +++ /dev/null @@ -1,724 +0,0 @@ -"""Executed by the PyO3 retained test with the actual built module in `native`.""" - -import asyncio -import contextvars -import gc -import inspect -import json -import threading -import unittest -import weakref -from copy import deepcopy -from datetime import datetime -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from types import ModuleType -from unittest.mock import patch - -import httpx - -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.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig -from litellm.rust_bridge.ocr_retained import OCREncoded, OCRRetainedBoundary, OCRRoots - -native: ModuleType = globals()["native"] -context = contextvars.ContextVar("retained-real-boundary", default="unset") -MODEL = "mistral-ocr-latest" -RESPONSE = {"pages": [{"index": 0, "markdown": "local OCR"}], "model": MODEL, "usage_info": {"pages_processed": 1}} - - -class Graph(dict): - pass - - -class Header(str): - pass - - -class PreCallAbort(BaseException): - pass - - -class CopiedDocumentBoundary(OCRRetainedBoundary): - def prepare(self) -> OCRRoots: - self.document = deepcopy(self.document) - return super().prepare() - - async def aprepare(self) -> OCRRoots: - self.document = deepcopy(self.document) - return await super().aprepare() - - -class ReboundBodyBoundary(OCRRetainedBoundary): - def encode(self, roots: OCRRoots) -> OCREncoded: - headers, url, _body, files = roots - view = self.logging_obj.model_call_details["additional_args"] - return super().encode((headers, url, view["complete_input_dict"], files)) - - -class ReboundHeadersBoundary(OCRRetainedBoundary): - def encode(self, roots: OCRRoots) -> OCREncoded: - _headers, url, body, files = roots - view = self.logging_obj.model_call_details["additional_args"] - return super().encode((view["headers"], url, body, files)) - - -class Callback(CustomLogger): - def __init__(self, action): - super().__init__() - self.action = action - self.failures = [] - self.calls = 0 - - def log_pre_api_call(self, model, messages, kwargs): - self.calls += 1 - try: - return self.action(kwargs["additional_args"]) - except AssertionError as error: - self.failures.append(str(error)) - raise - - -class Server: - def __init__(self): - self.requests = [] - self.started = threading.Event() - self.release = threading.Event() - self.finished = threading.Event() - owner = self - - class Handler(BaseHTTPRequestHandler): - def log_message(self, *args): - pass - - def do_POST(self): - body = self.rfile.read(int(self.headers["Content-Length"])) - owner.requests.append((self.path, sorted((k.lower(), v) for k, v in self.headers.items()), body)) - if self.path == "/blocked/v1/ocr": - owner.started.set() - if not owner.release.wait(10): - owner.finished.set() - return - failed = self.path == "/error/v1/ocr" - payload = b'{"message":"controlled HTTP failure"}' if failed else json.dumps(RESPONSE).encode() - self.send_response(429 if failed else 200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - try: - self.wfile.write(payload) - except (BrokenPipeError, ConnectionResetError): - pass - finally: - if self.path == "/blocked/v1/ocr": - owner.finished.set() - - self.httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) - self.thread.start() - self.url = f"http://127.0.0.1:{self.httpd.server_port}" - - def close(self): - self.release.set() - self.httpd.shutdown() - self.httpd.server_close() - self.thread.join(5) - assert not self.thread.is_alive() - - -def inputs(server, callbacks=(), *, document=None, optional=None, path="", client=None): - return { - "model": MODEL, - "document": Graph(type="document_url", document_url="https://example.test/original.pdf") - if document is None - else document, - "optional_params": {} if optional is None else optional, - "logging_obj": Logging( - model=MODEL, - messages=[], - stream=False, - call_type="ocr", - start_time=datetime.now(), - litellm_call_id="retained-real", - function_id="retained-real", - dynamic_input_callbacks=list(callbacks), - supports_correlation_logging=False, - ), - "api_key": "local-test-key", - "api_base": server.url + path, - "headers": {"X-Proof": "original"}, - "provider_config": MistralOCRConfig(), - "litellm_params": {}, - "custom_llm_provider": "mistral", - "timeout": 5.0, - "client": client, - } - - -def invoke(mode, kwargs, *, boundary_factory=OCRRetainedBoundary): - handler = BaseLLMHTTPHandler() - if mode == "python-sync": - return handler.ocr(**kwargs) - if mode == "python-async": - return handler.async_ocr(**kwargs) - if mode == "native-sync": - return native.ocr_retained(boundary_factory(handler=handler, **kwargs)) - assert mode == "native-async" - return native.aocr_retained(boundary_factory(handler=handler, **kwargs)) - - -class RealBoundaryTests(unittest.TestCase): - def setUp(self): - self.server = Server() - self.addCleanup(self.server.close) - self.sync_client = HTTPHandler(timeout=5.0) - self.addCleanup(self.sync_client.close) - - def check_callbacks(self, callbacks): - for callback in callbacks: - self.assertEqual(callback.failures, []) - self.assertEqual(callback.calls, 1) - - async def differential(self, mode, *, boundary_factory=OCRRetainedBoundary): - context.set("caller") - thread = threading.get_ident() - task = asyncio.current_task() - document = Graph(type="document_url", document_url="https://example.test/original.pdf") - nested = Graph(values=[1]) - optional = {"unknown_python_json": {7: ("tuple", 2)}, "nested": nested, "nested_alias": nested["values"]} - retained = {} - events = [] - - def phase(name, expected): - self.assertEqual(threading.get_ident(), thread) - self.assertIs(asyncio.current_task(), task) - self.assertEqual(context.get(), expected) - events.append(name) - - def mutate(view): - phase("mutate", "caller") - body, headers = view["complete_input_dict"], view["headers"] - self.assertIs(body["document"], document, "caller document identity was not retained") - self.assertIs(body["nested"], nested) - self.assertIs(body["nested_alias"], nested["values"]) - self.assertIs(body["unknown_python_json"], optional["unknown_python_json"]) - retained.update(body=body, headers=headers, view=view) - headers["X-Proof"] = "in-place" - body["body_mutation"] = True - nested["values"].append(2) - view["headers"] = {"X-Proof": "must-not-send"} - view["complete_input_dict"] = {"document": {"document_url": "must-not-send"}} - context.set("mutated") - child_callback = Callback(lambda _: events.append("reentry")) - child = invoke("native-sync", inputs(self.server, [child_callback], path="/child", client=self.sync_client)) - self.assertEqual(child.pages[0].markdown, "local OCR") - self.check_callbacks([child_callback]) - return {"headers": {"X-Proof": "ignored-return"}, "complete_input_dict": {"invalid": object()}} - - def mutate_then_raise(view): - phase("raise", "mutated") - self.assertIs(view, retained["view"]) - self.assertEqual(view["complete_input_dict"], {"document": {"document_url": "must-not-send"}}) - self.assertEqual(view["headers"], {"X-Proof": "must-not-send"}) - retained["body"]["before_error"] = True - retained["headers"]["X-Before-Error"] = "yes" - context.set("caught") - raise RuntimeError("intentional non-blocking pre_call error") - - def closure_only(view): - phase("later", "caught") - self.assertIs(view, retained["view"]) - self.assertIsNot(view["complete_input_dict"], retained["body"]) - self.assertIsNot(view["headers"], retained["headers"]) - self.assertTrue(retained["body"]["before_error"]) - self.assertEqual(retained["headers"]["X-Before-Error"], "yes") - self.assertIs(retained["body"]["nested_alias"], nested["values"]) - self.assertEqual(retained["body"]["nested_alias"], [1, 2]) - view["complete_input_dict"]["observed"] = True - view["headers"]["X-View-Only"] = "not-on-wire" - document["document_url"] = "https://example.test/closure.pdf" - context.set("later") - - callbacks = [Callback(action) for action in (mutate, mutate_then_raise, closure_only)] - async_client = AsyncHTTPHandler(timeout=5.0) - try: - kwargs = inputs( - self.server, - callbacks, - document=document, - optional=optional, - client=async_client if mode.endswith("async") else self.sync_client, - ) - logging_ref = weakref.ref(kwargs["logging_obj"]) - before = len(self.server.requests) - pending = invoke(mode, kwargs, boundary_factory=boundary_factory) - if mode.endswith("async"): - self.assertTrue(inspect.iscoroutine(pending)) - self.assertEqual(events, []) - self.assertEqual(len(self.server.requests), before) - self.assertNotIn("additional_args", kwargs["logging_obj"].model_call_details) - response = await pending - else: - response = pending - self.check_callbacks(callbacks) - self.assertEqual(events, ["mutate", "reentry", "raise", "later"]) - self.assertEqual(context.get(), "later") - self.assertEqual(len(self.server.requests), before + 2) - wire = self.server.requests[-1] - self.assertEqual(wire[0], "/v1/ocr") - expected = ( - b'{"model":"mistral-ocr-latest","document":{"type":"document_url",' - b'"document_url":"https://example.test/closure.pdf"},"unknown_python_json":{"7":["tuple",2]},' - b'"nested":{"values":[1,2]},"nested_alias":[1,2],"body_mutation":true,"before_error":true}' - ) - self.assertEqual(wire[2], expected, "wire body must encode retained mutations after pre_call") - self.assertIn(("x-proof", "in-place"), wire[1], "wire headers must use retained execution headers") - self.assertIn(("x-before-error", "yes"), wire[1]) - self.assertIn(("authorization", "Bearer local-test-key"), wire[1]) - self.assertFalse(any(name == "x-view-only" for name, _ in wire[1])) - self.assertEqual( - retained["view"]["complete_input_dict"], - {"document": {"document_url": "must-not-send"}, "observed": True}, - ) - self.assertEqual(retained["view"]["headers"], {"X-Proof": "must-not-send", "X-View-Only": "not-on-wire"}) - del kwargs, pending - gc.collect() - self.assertIsNone(logging_ref(), "logging owner survived the completed call") - self.assertIs(retained["body"]["document"], document) - retained["body"]["nested"]["values"].append(3) - retained["headers"]["X-After-Return"] = "usable" - self.assertEqual(optional["nested"]["values"], [1, 2, 3]) - self.assertIs(retained["body"]["nested_alias"], optional["nested"]["values"]) - self.assertEqual(retained["body"]["nested_alias"], [1, 2, 3]) - self.assertEqual(retained["headers"]["X-After-Return"], "usable") - self.assertEqual(wire[2], expected) - self.assertEqual(response.pages[0].markdown, "local OCR") - return wire, response.model_dump() - finally: - await async_client.close() - - def test_differential_callbacks_wire(self): - async def exercise(): - baseline = await self.differential("python-sync") - for mode in ("python-async", "native-sync", "native-async"): - with self.subTest(mode=mode): - self.assertEqual(await self.differential(mode), baseline) - - asyncio.run(exercise()) - - @unittest.expectedFailure - def test_known_gap_send_time_auth_and_request_hook(self): - """UC-HTTPX-SEND: pre_call, client auth, then request hook must determine the provider's wire headers.""" - observations = [] - for mode in ("python-sync", "native-sync"): - events = [] - - def request_hook(request, events=events): - events.append(("request_hook", request.headers["Authorization"])) - request.headers["X-Send-Hook"] = "present" - - callback = Callback( - lambda view, events=events: events.append(("pre_call", view["headers"]["Authorization"])) - ) - with httpx.Client(auth=("local-user", "local-password"), event_hooks={"request": [request_hook]}) as client: - before = len(self.server.requests) - response = invoke(mode, inputs(self.server, [callback], client=HTTPHandler(client=client))) - self.check_callbacks([callback]) - self.assertEqual(response.pages[0].markdown, "local OCR") - self.assertEqual(len(self.server.requests), before + 1) - headers = dict(self.server.requests[-1][1]) - observations.append((events, headers["authorization"], headers.get("x-send-hook"))) - baseline, retained = observations - self.assertEqual(baseline[0][0], ("pre_call", "Bearer local-test-key")) - self.assertEqual(baseline[0][1], ("request_hook", baseline[1])) - self.assertTrue(baseline[1].startswith("Basic ")) - self.assertEqual(baseline[2], "present") - self.assertEqual(retained, baseline, "UC-HTTPX-SEND: retained send omitted client auth/request hook") - - @unittest.expectedFailure - def test_known_gap_custom_transport(self): - """UC-HTTPX-SEND: after pre_call, the client's transport must select the response without a network POST.""" - observations = [] - for mode in ("python-sync", "native-sync"): - events = [] - - def transport(request, events=events): - events.append("transport") - return httpx.Response(200, json={**RESPONSE, "pages": [{"index": 0, "markdown": "custom transport"}]}) - - callback = Callback(lambda view, events=events: events.append("pre_call")) - with httpx.Client(transport=httpx.MockTransport(transport)) as client: - before = len(self.server.requests) - response = invoke(mode, inputs(self.server, [callback], client=HTTPHandler(client=client))) - self.check_callbacks([callback]) - observations.append((events, len(self.server.requests) - before, response.pages[0].markdown)) - baseline, retained = observations - self.assertEqual(baseline, (["pre_call", "transport"], 0, "custom transport")) - self.assertEqual(retained, baseline, "UC-HTTPX-SEND: retained send bypassed the custom transport") - - @unittest.expectedFailure - def test_known_gap_pre_call_timeout_mutation(self): - """UC-TIMEOUT-MUTATION: a caller's timeout is changed by pre_call; encoding must accept read=None.""" - observations = [] - for mode in ("python-sync", "native-sync"): - timeout = httpx.Timeout(5.0) - - def mutate_timeout(view, timeout=timeout): - timeout.read = None - - callback = Callback(mutate_timeout) - kwargs = inputs(self.server, [callback], client=self.sync_client) - kwargs["timeout"] = timeout - before = len(self.server.requests) - try: - response = invoke(mode, kwargs) - except ValueError as error: - outcome = ("error", str(error)) - else: - outcome = ("success", response.pages[0].markdown) - self.check_callbacks([callback]) - self.assertIsNone(timeout.read) - observations.append((outcome, len(self.server.requests) - before)) - baseline, retained = observations - self.assertEqual(baseline, (("success", "local OCR"), 1)) - self.assertEqual(retained, baseline, "UC-TIMEOUT-MUTATION: retained encoding rejected the callback's timeout") - - def check_negative_control(self, boundary_factory, failure, expected_requests): - async def exercise(): - baseline = await self.differential("python-sync") - for mode in ("native-sync", "native-async"): - with self.subTest(mode=mode): - symbol = "aocr_retained" if mode.endswith("async") else "ocr_retained" - self.assertTrue(inspect.isbuiltin(getattr(native, symbol))) - self.assertEqual(await self.differential(mode), baseline) - before = len(self.server.requests) - with self.assertRaisesRegex(AssertionError, failure): - await self.differential(mode, boundary_factory=boundary_factory) - self.assertEqual(len(self.server.requests), before + expected_requests) - self.assertEqual(self.server.requests[-1][0], "/v1/ocr") - - asyncio.run(exercise()) - - def test_negative_control_copied_caller_document(self): - self.check_negative_control(CopiedDocumentBoundary, "caller document identity was not retained", 1) - - def test_negative_control_rebound_logging_body(self): - self.check_negative_control(ReboundBodyBoundary, "wire body must encode retained mutations after pre_call", 2) - - def test_negative_control_rebound_logging_headers(self): - self.check_negative_control(ReboundHeadersBoundary, "wire headers must use retained execution headers", 2) - - def test_differential_callback_retained_mutation_after_post_received(self): - async def suspended(mode): - self.server.started.clear() - self.server.release.clear() - self.server.finished.clear() - retained = {} - - def retain(view): - retained.update(body=view["complete_input_dict"], headers=view["headers"], view=view) - view["complete_input_dict"] = {"replacement": True} - view["headers"] = {"X-Proof": "must-not-send"} - - callback = Callback(retain) - client = AsyncHTTPHandler(timeout=5.0) - kwargs = inputs(self.server, [callback], path="/blocked", client=client) - logging = kwargs["logging_obj"] - logging.log_raw_request_response = True - before = len(self.server.requests) - task = asyncio.create_task(invoke(mode, kwargs)) - try: - self.assertTrue(await asyncio.to_thread(self.server.started.wait, 5)) - self.assertFalse(task.done()) - self.assertFalse(self.server.release.is_set()) - self.assertFalse(self.server.finished.is_set()) - self.check_callbacks([callback]) - self.assertEqual(len(self.server.requests), before + 1) - path, headers, body = self.server.requests[-1] - wire = (path, tuple(headers), body) - expected = ( - b'{"model":"mistral-ocr-latest","document":{"type":"document_url",' - b'"document_url":"https://example.test/original.pdf"}}' - ) - self.assertEqual(path, "/blocked/v1/ocr") - self.assertEqual(body, expected) - self.assertIn(("x-proof", "original"), headers) - self.assertNotIn(("x-proof", "must-not-send"), headers) - logged_body = logging.model_call_details["raw_request_typed_dict"]["raw_request_body"] - self.assertIs(logged_body, retained["body"]) - self.assertIs(logged_body["document"], kwargs["document"]) - self.assertIs(logging.model_call_details["additional_args"], retained["view"]) - self.assertIsNot(retained["view"]["complete_input_dict"], retained["body"]) - self.assertIsNot(retained["view"]["headers"], retained["headers"]) - - retained["body"]["after_encoding"] = True - retained["body"]["document"]["document_url"] = "https://example.test/while-blocked.pdf" - retained["headers"]["X-Proof"] = "while-blocked" - self.assertTrue(logged_body["after_encoding"]) - self.assertEqual(kwargs["document"]["document_url"], "https://example.test/while-blocked.pdf") - self.assertEqual(logged_body["document"]["document_url"], "https://example.test/while-blocked.pdf") - self.assertEqual(retained["headers"]["X-Proof"], "while-blocked") - self.assertEqual(retained["view"]["complete_input_dict"], {"replacement": True}) - self.assertEqual(retained["view"]["headers"], {"X-Proof": "must-not-send"}) - self.assertFalse(task.done()) - self.assertFalse(self.server.finished.is_set()) - self.assertEqual((path, tuple(headers), body), wire) - - self.server.release.set() - response = await asyncio.wait_for(task, 5) - self.assertEqual(response.pages[0].markdown, "local OCR") - self.assertIs(logging.model_call_details["raw_request_typed_dict"]["raw_request_body"], logged_body) - self.assertTrue(logged_body["after_encoding"]) - self.assertEqual(len(self.server.requests), before + 1) - received_path, received_headers, received_body = self.server.requests[-1] - self.assertEqual((received_path, tuple(received_headers), received_body), wire) - self.check_callbacks([callback]) - return wire, logged_body, retained["headers"], response.model_dump() - finally: - self.server.release.set() - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - await client.close() - if self.server.started.is_set(): - self.assertTrue(await asyncio.to_thread(self.server.finished.wait, 5)) - - async def exercise(): - baseline = await suspended("python-async") - self.assertEqual(await suspended("native-async"), baseline) - - asyncio.run(exercise()) - - def test_public_rust_dispatch_wire_fallback_and_escaping_base_exception(self): - async def exercise(): - for asynchronous in (False, True): - for outcome in ("success", "missing-symbol", "pre-call-abort", "disabled"): - with self.subTest(asynchronous=asynchronous, outcome=outcome): - symbol = "aocr_retained" if asynchronous else "ocr_retained" - self.assertTrue(inspect.isbuiltin(getattr(native, symbol))) - missing_symbol = ModuleType("native_without_" + symbol) - missing_symbol.__dict__.update( - (name, value) for name, value in vars(native).items() if name != symbol - ) - escaped = PreCallAbort("public pre_call must escape unchanged") - before = len(self.server.requests) - - def mutate(view, outcome=outcome, escaped=escaped): - view["headers"]["X-Proof"] = "public-in-place" - view["complete_input_dict"]["public_mutation"] = True - view["complete_input_dict"]["document"]["document_url"] = ( - "https://example.test/public-mutated.pdf" - ) - if outcome == "pre-call-abort": - raise escaped - - callback = Callback(mutate) - kwargs = inputs(self.server, [callback]) - with patch( - "litellm.rust_bridge.get_native_bridge", - return_value=missing_symbol if outcome == "missing-symbol" else native, - ) as loader: - public_kwargs = { - **{ - key: kwargs[key] - for key in ( - "model", - "document", - "api_key", - "api_base", - "custom_llm_provider", - "timeout", - ) - }, - "extra_headers": kwargs["headers"], - "litellm_logging_obj": kwargs["logging_obj"], - "rust": outcome != "disabled", - } - - async def call(asynchronous=asynchronous, public_kwargs=public_kwargs): - if asynchronous: - return await litellm.aocr(**public_kwargs) - return litellm.ocr(**public_kwargs) - - if outcome == "pre-call-abort": - with self.assertRaises(PreCallAbort) as caught: - await call() - self.assertIs(caught.exception, escaped) - else: - response = await call() - self.assertEqual(response.pages[0].markdown, "local OCR") - - if outcome == "disabled": - loader.assert_not_called() - else: - loader.assert_called_once_with() - self.check_callbacks([callback]) - self.assertEqual(len(self.server.requests), before + int(outcome != "pre-call-abort")) - if outcome != "pre-call-abort": - path, headers, body = self.server.requests[-1] - self.assertEqual(path, "/v1/ocr") - self.assertIn(("x-proof", "public-in-place"), headers) - self.assertIn(("authorization", "Bearer local-test-key"), headers) - self.assertEqual( - json.loads(body), - { - "model": MODEL, - "document": { - "type": "document_url", - "document_url": "https://example.test/public-mutated.pdf", - }, - "public_mutation": True, - }, - ) - - asyncio.run(exercise()) - - def lifecycle_inputs(self, outcome, retained): - refs = [] - - def callback(view): - body, headers = view["complete_input_dict"], view["headers"] - body["sentinel"] = Graph(alive=True) - headers["X-Sentinel"] = Header("alive") - refs.extend((weakref.ref(body["sentinel"]), weakref.ref(headers["X-Sentinel"]))) - if outcome == "encoding": - body["not_json"] = object() - if retained is not None: - retained.extend((body, headers)) - view["complete_input_dict"] = {} - view["headers"] = {} - if outcome == "pre-call-abort": - raise PreCallAbort("lifecycle pre_call abort") - - logger = Callback(callback) - kwargs = inputs( - self.server, - [logger], - optional={"nested": Graph(alive=True)}, - path={"http": "/error", "cancel": "/blocked"}.get(outcome, ""), - client=self.sync_client, - ) - refs.extend(weakref.ref(kwargs[key]) for key in ("document", "logging_obj")) - refs.append(weakref.ref(kwargs["optional_params"]["nested"])) - return kwargs, logger, refs - - async def lifecycle(self, mode, outcome, retained=None): - kwargs, logger, refs = self.lifecycle_inputs(outcome, retained) - before = len(self.server.requests) - async_client = AsyncHTTPHandler(timeout=5.0) - if mode.endswith("async"): - kwargs["client"] = async_client - try: - try: - pending = invoke(mode, kwargs) - response = await pending if mode.endswith("async") else pending - except BaseLLMException as error: - self.assertIn(outcome, ("encoding", "http")) - self.assertEqual(error.status_code, 500 if outcome == "encoding" else 429) - signature = (type(error), error.status_code, str(error)) - except PreCallAbort as error: - self.assertEqual(outcome, "pre-call-abort") - signature = (type(error), str(error)) - else: - self.assertEqual(outcome, "success") - self.assertEqual(response.pages[0].markdown, "local OCR") - signature = response.model_dump() - self.check_callbacks([logger]) - self.assertEqual(len(refs), 5) - self.assertEqual(len(self.server.requests), before + (outcome not in ("encoding", "pre-call-abort"))) - return refs, signature - finally: - await async_client.close() - - def test_collection_after_success_and_failures(self): - async def exercise(): - for outcome in ("success", "encoding", "http", "pre-call-abort"): - baseline = None - for mode in ("python-sync", "python-async", "native-sync", "native-async"): - with self.subTest(mode=mode, outcome=outcome): - refs, signature = await self.lifecycle(mode, outcome) - gc.collect() - self.assertTrue(all(ref() is None for ref in refs), f"request graph leaked: {mode=} {outcome=}") - if baseline is None: - baseline = signature - self.assertEqual(signature, baseline) - - asyncio.run(exercise()) - - def test_callback_retained_graph_remains_usable_then_collects(self): - async def exercise(): - for outcome in ("success", "encoding", "http", "pre-call-abort"): - for mode in ("python-sync", "python-async", "native-sync", "native-async"): - with self.subTest(mode=mode, outcome=outcome): - retained = [] - refs, _ = await self.lifecycle(mode, outcome, retained) - gc.collect() - self.assertIsNone(refs[1](), f"logging owner survived: {mode=} {outcome=}") - self.assertTrue(all(refs[index]() is not None for index in (0, 2, 3, 4))) - self.assertIs(retained[0]["document"], refs[0]()) - self.assertIs(retained[0]["nested"], refs[2]()) - self.assertIs(retained[0]["sentinel"], refs[3]()) - self.assertIs(retained[1]["X-Sentinel"], refs[4]()) - retained[0]["document"]["after_return"] = "usable" - retained[0]["nested"]["alive"] = "nested still usable" - retained[0]["sentinel"]["alive"] = "still usable" - retained[1]["X-After-Return"] = "usable" - self.assertEqual(refs[0]()["after_return"], "usable") - self.assertEqual(refs[2]()["alive"], "nested still usable") - self.assertEqual(refs[3]()["alive"], "still usable") - self.assertEqual(retained[1]["X-After-Return"], "usable") - retained.clear() - gc.collect() - self.assertTrue(all(ref() is None for ref in refs)) - - asyncio.run(exercise()) - - def test_collection_after_cancellation_during_blocked_transport(self): - async def cancel(): - kwargs, logger, refs = self.lifecycle_inputs("cancel", None) - client = AsyncHTTPHandler(timeout=5.0) - kwargs["client"] = client - task = asyncio.create_task(invoke("native-async", kwargs)) - try: - self.assertTrue(await asyncio.to_thread(self.server.started.wait, 5)) - self.check_callbacks([logger]) - self.assertEqual(len(refs), 5) - gc.collect() - self.assertTrue(all(ref() is not None for ref in refs)) - task.cancel() - with self.assertRaises(asyncio.CancelledError): - await task - self.assertFalse(self.server.release.is_set()) - return refs - finally: - if not task.done(): - task.cancel() - await asyncio.gather(task, return_exceptions=True) - await client.close() - - async def exercise(): - refs = await cancel() - barrier = asyncio.get_running_loop().create_future() - asyncio.get_running_loop().call_soon(barrier.set_result, None) - await barrier - gc.collect() - self.assertTrue(all(ref() is None for ref in refs), "cancelled native call retained the request graph") - self.server.release.set() - self.assertTrue(await asyncio.to_thread(self.server.finished.wait, 5)) - self.assertEqual(len(self.server.requests), 1) - - asyncio.run(exercise()) - - -def run_case(scenario): - suite = unittest.TestSuite([RealBoundaryTests("test_" + scenario)]) - result = unittest.TextTestRunner(verbosity=2).run(suite) - for _, traceback in result.expectedFailures: - assert "AssertionError:" in traceback and ( - "UC-HTTPX-SEND:" in traceback or "UC-TIMEOUT-MUTATION:" in traceback - ), traceback - assert result.wasSuccessful(), "OCR retained parity test failed or unexpectedly succeeded" diff --git a/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py b/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py deleted file mode 100644 index f94db3386f0..00000000000 --- a/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py +++ /dev/null @@ -1,663 +0,0 @@ -import asyncio -import contextvars -import gc -import http.client -import http.server -import inspect -import sys -import threading -import weakref -from copy import deepcopy -from typing import NamedTuple - -native = globals()["native"] -started = threading.Event() -release = threading.Event() -server = None -server_thread = None -url = None -context = contextvars.ContextVar("proof", default="unset") - - -def start_server(): - global server, server_thread, url - requests = [] - - class Handler(http.server.BaseHTTPRequestHandler): - def log_message(self, *args): - pass - - def do_POST(self): - body = self.rfile.read(int(self.headers["Content-Length"])) - requests.append((self.path, self.headers.get_all("X-Proof"), body)) - if body == b"hold": - started.set() - release.wait(5) - self.send_response(429) - self.send_header("X-Reply", "one") - self.send_header("X-Reply", "two") - self.send_header("Content-Length", "3") - self.end_headers() - try: - self.wfile.write(b"\x00\xffR") - except (BrokenPipeError, ConnectionResetError): - pass - - server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - url = "http://127.0.0.1:%s/ocr" % server.server_port - return requests - - -requests = start_server() - - -def stop_server(): - release.set() - server.shutdown() - server.server_close() - server_thread.join(5) - assert not server_thread.is_alive() - - -class Graph(dict): - pass - - -class Boundary: - def __init__(self, *, asynchronous=False, nested=False, failure=None, hold=False): - self.asynchronous = asynchronous - self.nested = nested - self.failure = failure - self.hold = hold - self.events = [] - self.thread = threading.get_ident() - self.task = asyncio.current_task() if asynchronous else None - self.result = object() - self.error = LookupError("original callback error") - - def phase(self, name): - assert threading.get_ident() == self.thread - if self.asynchronous: - assert asyncio.current_task() is self.task - assert context.get() == ("initial" if name == "prepare" else "prepared") - self.events.append(name) - if self.failure == name: - raise self.error - if not self.nested and not self.hold: - child = Boundary(nested=True) - assert native.ocr_retained(child) is child.result - assert child.events == ["prepare", "encode", "finish"] - - def prepare(self): - self.phase("prepare") - headers = Graph({"X-Proof": "original"}) - document = object() - body = Graph(document=document, alias=document) - body["cycle"] = body - self.refs = (weakref.ref(headers), weakref.ref(body)) - self.view = {"headers": headers, "body": body} - headers["X-Proof"] = "mutated" - self.view["headers"] = {"replacement": True} - self.view["body"] = {"replacement": True} - return (headers, url, body, None) - - async def aprepare(self): - await asyncio.sleep(0) - roots = self.prepare() - context.set("prepared") - return roots - - def encode(self, roots): - self.phase("encode") - headers, target, body, files = roots - assert headers is self.refs[0]() and body is self.refs[1]() - assert headers["X-Proof"] == "mutated" - assert body["document"] is body["alias"] and body["cycle"] is body - assert files is None - assert self.view == {"headers": {"replacement": True}, "body": {"replacement": True}} - return ( - target, - [(b"X-Proof", b"mutated"), (b"X-Proof", b"duplicate")], - b"hold" if self.hold else b"\x00\xffQ", - 3.0, - ) - - def finish(self, wire): - self.phase("finish") - assert type(wire) is tuple and len(wire) == 3 - status, headers, content = wire - assert status == 429 - assert type(headers) is list - assert all(type(pair) is tuple and all(type(v) is bytes for v in pair) for pair in headers) - assert [v for k, v in headers if k == b"x-reply"] == [b"one", b"two"] - assert type(content) is bytes and content == b"\x00\xffR" - assert all(ref() is not None for ref in self.refs) - return self.result - - async def afinish(self, wire): - await asyncio.sleep(0) - return self.finish(wire) - - -def collected(boundary): - gc.collect() - assert all(ref() is None for ref in boundary.refs) - assert boundary.view == {"headers": {"replacement": True}, "body": {"replacement": True}} - - -def run_cold_cache_reentry(filename): - """UC-COLD-REENTRY: cold compilation permits nested native calls and subsequent transport.""" - events = [] - error = LookupError("cold-cache preparation error") - - class FailingBoundary: - async def aprepare(self): - raise error - - def invoke_failure(): - pending = native.aocr_retained(FailingBoundary()) - observed = None - try: - pending.send(None) - except LookupError as caught: - observed = caught - finally: - pending.close() - error.__traceback__ = None - assert observed is error - - def audit(event, args): - if event == "compile" and args[1] == filename and not events: - events.append("entered") - print(f"UC-COLD-REENTRY: entering {filename}", flush=True) # noqa: T201 # diagnose a deadlocked child process - invoke_failure() - events.append("nested completed") - - async def successful_async(): - context.set("initial") - boundary = Boundary(asynchronous=True, nested=True) - assert await native.aocr_retained(boundary) is boundary.result - assert boundary.events == ["prepare", "encode", "finish"] - collected(boundary) - - try: - if filename == "retained_callback.py": - native.aocr_retained(object()).close() - sys.addaudithook(audit) - invoke_failure() - events.append("outer completed") - assert events == ["entered", "nested completed", "outer completed"] - assert requests == [] - boundary = Boundary(nested=True) - assert native.ocr_retained(boundary) is boundary.result - assert boundary.events == ["prepare", "encode", "finish"] - collected(boundary) - asyncio.run(successful_async()) - assert len(requests) == 2 - finally: - stop_server() - - -def check_error(boundary, error, phase): - assert error is boundary.error - names = [] - traceback = error.__traceback__ - while traceback: - names.append(traceback.tb_frame.f_code.co_name) - traceback = traceback.tb_next - assert phase in names and "phase" in names - assert boundary.events == ["prepare", "encode", "finish"][: ["prepare", "encode", "finish"].index(phase) + 1] - - -async def exercise(): - context.set("initial") - boundary = Boundary(asynchronous=True) - pending = native.aocr_retained(boundary) - assert inspect.iscoroutine(pending) - assert boundary.events == [] - assert await pending is boundary.result - assert context.get() == "prepared" - assert boundary.events == ["prepare", "encode", "finish"] - collected(boundary) - - unused = Boundary(asynchronous=True) - ref = weakref.ref(unused) - pending = native.aocr_retained(unused) - assert unused.events == [] - del unused - assert ref() is not None - pending.close() - del pending - gc.collect() - assert ref() is None - - for phase in ("prepare", "encode", "finish"): - context.set("initial") - boundary = Boundary(asynchronous=True, nested=True, failure=phase) - try: - await native.aocr_retained(boundary) - except LookupError as error: - check_error(boundary, error, phase) - else: - raise AssertionError("callback error was swallowed") - boundary.error.__traceback__ = None - if phase != "prepare": - collected(boundary) - - context.set("initial") - boundary = Boundary(asynchronous=True, hold=True) - - async def cancellable(): - boundary.task = asyncio.current_task() - await native.aocr_retained(boundary) - - task = asyncio.create_task(cancellable()) - assert await asyncio.to_thread(started.wait, 2) - gc.collect() - assert all(ref() is not None for ref in boundary.refs) - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - else: - raise AssertionError("cancellation was swallowed") - assert boundary.events == ["prepare", "encode"] - del task - await asyncio.sleep(0) - collected(boundary) - release.set() - - -def run_ownership_contract(): - try: - boundary = Boundary() - assert native.ocr_retained(boundary) is boundary.result - assert boundary.events == ["prepare", "encode", "finish"] - collected(boundary) - for phase in ("prepare", "encode", "finish"): - boundary = Boundary(nested=True, failure=phase) - try: - native.ocr_retained(boundary) - except LookupError as error: - check_error(boundary, error, phase) - else: - raise AssertionError("callback error was swallowed") - boundary.error.__traceback__ = None - if phase != "prepare": - collected(boundary) - boundary = Boundary(hold=True) - observations = [] - - def observe_blocked_transport(): - received = started.wait(2) - gc.collect() - observations.append((received, tuple(ref() is not None for ref in boundary.refs))) - release.set() - - observer = threading.Thread(target=observe_blocked_transport) - observer.start() - try: - assert native.ocr_retained(boundary) is boundary.result - finally: - release.set() - observer.join(5) - assert not observer.is_alive() - assert observations == [(True, (True, True))] - assert boundary.events == ["prepare", "encode", "finish"] - collected(boundary) - started.clear() - release.clear() - asyncio.run(asyncio.wait_for(exercise(), 15)) - assert requests - assert all( - path == "/ocr" and headers == ["mutated", "duplicate"] and body in (b"\x00\xffQ", b"hold") - for path, headers, body in requests - ) - finally: - stop_server() - - -class TimeoutBoundary(Boundary): - def __init__(self, *, timeout, url, asynchronous=False): - super().__init__(asynchronous=asynchronous) - self.timeout = timeout - self.url = url - - def prepare(self): - return ({}, self.url, {}, None) - - async def aprepare(self): - return self.prepare() - - def encode(self, roots): - headers, target, body, files = roots - return (target, [], b"hold", self.timeout) - - def finish(self, wire): - raise AssertionError("client-side failure must not reach finish") - - -def run_error_contract(): - cases = [ - ("timeout", url, 0.05), - ("refused", "http://127.0.0.1:1/", 1.0), - ] - try: - for name, target, timeout in cases: - boundary = TimeoutBoundary(timeout=timeout, url=target) - try: - native.ocr_retained(boundary) - except RuntimeError: - pass - else: - raise AssertionError(f"{name} did not surface as RuntimeError") - finally: - stop_server() - - -class FrozenObservation(NamedTuple): - root_tuple_identity: bool - roots_identity: tuple[bool, bool] - document_is_caller: bool - document_alias: bool - logging_envelope_identity: bool - logging_roots_identity: tuple[bool, bool] - logging_fields: tuple[str, str] - retained_fields: tuple[str, str] - caller_value: str - caught_error_observation: tuple[str, ...] - callback_count: int - phases: tuple[str, ...] - request: tuple[str, tuple[str, ...], bytes] - response: tuple[int, tuple[bytes, ...], bytes] - - -class ContractCallback: - def __init__(self, scenario, caller): - self.scenario = scenario - self.caller = caller - self.mutate_caller = lambda: caller.__setitem__("value", "closure") - self.calls = 0 - self.error = LookupError("caught callback failure") - self.writer = None - self.write = threading.Event() - self.written = threading.Event() - - def __call__(self, view): - self.calls += 1 - self.view = view - self.headers, self.body = view["headers"], view["body"] - if self.scenario == "caller_closure": - self.mutate_caller() - elif self.scenario == "field_replace": - document = Graph(value="replacement") - view["headers"] = Graph({"X-Proof": "replacement"}) - view["body"] = Graph(document=document, alias=document) - elif self.scenario == "delayed_after_prepare_before_encode_read": - - def writer(): - if self.write.wait(5): - self.mutate("delayed") - self.written.set() - - self.writer = threading.Thread(target=writer) - self.writer.start() - else: - assert self.scenario in ("retain_mutate", "mutate_then_caught_error_observe") - self.mutate("mutated" if self.scenario == "retain_mutate" else "caught") - if self.scenario == "mutate_then_caught_error_observe": - raise self.error - - def mutate(self, value): - self.headers["X-Proof"] = value - self.body["document"]["value"] = value - - def before_encode_read(self): - if self.writer is not None: - self.write.set() - assert self.written.wait(2), "scheduled mutation did not complete before encoder read" - self.writer.join(2) - assert not self.writer.is_alive() - - def close(self): - self.write.set() - if self.writer is not None: - self.writer.join(5) - assert not self.writer.is_alive() - self.error.__traceback__ = None - - -class TransportContractBinding: - def __init__(self, caller, callback): - self.caller = caller - self.callback = callback - self.phases = [] - self.caught = () - - def prepare(self): - self.phases.append("prepare") - headers = Graph({"X-Proof": "original"}) - body = Graph(document=self.caller, alias=self.caller) - self.view = {"headers": headers, "body": body} - self.prepared_roots = (headers, url, body, None) - try: - self.callback(self.view) - except LookupError as error: - if error is not self.callback.error: - raise - self.caught = (str(error), headers["X-Proof"], body["document"]["value"]) - return self.prepared_roots - - async def aprepare(self): - await asyncio.sleep(0) - return self.prepare() - - def encode(self, roots): - self.callback.before_encode_read() - self.phases.append("encode") - headers, target, body, files = roots - assert files is None - self.root_tuple_identity = roots is self.prepared_roots - self.roots_identity = (headers is self.callback.headers, body is self.callback.body) - self.document_is_caller = body["document"] is self.callback.caller - self.document_alias = body["document"] is body["alias"] - return ( - target, - [(b"X-Proof", headers["X-Proof"].encode())], - b"\x00" + body["document"]["value"].encode() + b"\xff", - 3.0, - ) - - def finish(self, wire): - self.phases.append("finish") - status, headers, content = wire - return (status, tuple(value for key, value in headers if key == b"x-reply"), content) - - async def afinish(self, wire): - await asyncio.sleep(0) - return self.finish(wire) - - -class BindingVariant: - def __init__(self, binding, variant): - self.binding = binding - self.variant = variant - - def before_prepare(self): - if self.variant == "copy_caller_inputs_before_prepare": - self.binding.caller = deepcopy(self.binding.caller) - - def at_prepare_return(self, roots): - if self.variant == "copy_roots_after_prepare": - return deepcopy(roots) - if self.variant == "reconstruct_root_tuple": - return tuple(list(roots)) - if self.variant == "reconstruct_logging_envelope": - self.binding.view = dict(self.binding.view) - return roots - - def prepare(self): - self.before_prepare() - return self.at_prepare_return(self.binding.prepare()) - - async def aprepare(self): - self.before_prepare() - return self.at_prepare_return(await self.binding.aprepare()) - - def encode(self, roots): - if self.variant == "encode_logging_replacements": - headers, target, body, files = roots - return self.binding.encode((self.binding.view["headers"], target, self.binding.view["body"], files)) - return self.binding.encode(roots) - - def finish(self, wire): - return self.binding.finish(wire) - - async def afinish(self, wire): - return await self.binding.afinish(wire) - - -def reference_post(encoded): - target, headers, body, timeout = encoded - assert target == url - connection = http.client.HTTPConnection("127.0.0.1", server.server_port, timeout=timeout) - try: - connection.putrequest("POST", "/ocr") - for key, value in headers: - connection.putheader(key.decode("ascii"), value.decode("ascii")) - connection.putheader("Content-Length", str(len(body))) - connection.endheaders(body) - response = connection.getresponse() - return ( - response.status, - [(key.lower().encode(), value.encode()) for key, value in response.getheaders()], - response.read(), - ) - finally: - connection.close() - - -async def reference_async(binding): - roots = await binding.aprepare() - wire = await asyncio.to_thread(reference_post, binding.encode(roots)) - return await binding.afinish(wire) - - -def original_observation(value, *, header=None, logging=None, caught=()): - fields = (header or value, value) - return FrozenObservation( - True, - (True, True), - True, - True, - True, - (False, False) if logging else (True, True), - logging or fields, - fields, - value, - caught, - 1, - ("prepare", "encode", "finish"), - ("/ocr", (fields[0],), b"\x00" + value.encode() + b"\xff"), - (429, (b"one", b"two"), b"\x00\xffR"), - ) - - -ORIGINAL_OBSERVATIONS = { - "retain_mutate": original_observation("mutated"), - "caller_closure": original_observation("closure", header="original"), - "delayed_after_prepare_before_encode_read": original_observation("delayed"), - "field_replace": original_observation("original", logging=("replacement", "replacement")), - "mutate_then_caught_error_observe": original_observation( - "caught", caught=("caught callback failure", "caught", "caught") - ), -} - -VARIANT_OBSERVATIONS = { - ("retain_mutate", "copy_caller_inputs_before_prepare"): { - "caller_value": "original", - "document_is_caller": False, - }, - ("caller_closure", "copy_caller_inputs_before_prepare"): { - "document_is_caller": False, - "logging_fields": ("original", "original"), - "retained_fields": ("original", "original"), - "request": ("/ocr", ("original",), b"\x00original\xff"), - }, - ("retain_mutate", "copy_roots_after_prepare"): { - "root_tuple_identity": False, - "roots_identity": (False, False), - "document_is_caller": False, - }, - ("delayed_after_prepare_before_encode_read", "copy_roots_after_prepare"): { - "root_tuple_identity": False, - "roots_identity": (False, False), - "document_is_caller": False, - "request": ("/ocr", ("original",), b"\x00original\xff"), - }, - ("mutate_then_caught_error_observe", "copy_roots_after_prepare"): { - "root_tuple_identity": False, - "roots_identity": (False, False), - "document_is_caller": False, - }, - ("field_replace", "encode_logging_replacements"): { - "root_tuple_identity": False, - "roots_identity": (False, False), - "document_is_caller": False, - "request": ("/ocr", ("replacement",), b"\x00replacement\xff"), - }, - ("retain_mutate", "reconstruct_logging_envelope"): {"logging_envelope_identity": False}, - ("field_replace", "reconstruct_logging_envelope"): {"logging_envelope_identity": False}, - ("retain_mutate", "reconstruct_root_tuple"): {"root_tuple_identity": False}, - ("delayed_after_prepare_before_encode_read", "reconstruct_root_tuple"): {"root_tuple_identity": False}, - ("field_replace", "reconstruct_root_tuple"): {"root_tuple_identity": False}, -} - - -def run_comparison_contract(scenario, variant): - try: - expected = ORIGINAL_OBSERVATIONS[scenario] - if variant != "original": - expected = expected._replace(**VARIANT_OBSERVATIONS[(scenario, variant)]) - for mode in ("reference-sync", "reference-async", "native-sync", "native-async"): - caller = Graph(value="original") - callback = ContractCallback(scenario, caller) - binding = TransportContractBinding(caller, callback) - selected = binding if variant == "original" else BindingVariant(binding, variant) - offset = len(requests) - try: - if mode == "reference-sync": - roots = selected.prepare() - response = selected.finish(reference_post(selected.encode(roots))) - elif mode == "reference-async": - response = asyncio.run(reference_async(selected)) - elif mode == "native-sync": - response = native.ocr_retained(selected) - else: - response = asyncio.run(native.aocr_retained(selected)) - assert len(requests) == offset + 1, (scenario, variant, mode, requests[offset:]) - path, headers, body = requests[offset] - observed = FrozenObservation( - binding.root_tuple_identity, - binding.roots_identity, - binding.document_is_caller, - binding.document_alias, - binding.view is callback.view, - (binding.view["headers"] is callback.headers, binding.view["body"] is callback.body), - (binding.view["headers"]["X-Proof"], binding.view["body"]["document"]["value"]), - (callback.headers["X-Proof"], callback.body["document"]["value"]), - caller["value"], - binding.caught, - callback.calls, - tuple(binding.phases), - (path, tuple(headers), body), - response, - ) - assert observed == expected, (scenario, variant, mode, observed, expected) - finally: - callback.close() - finally: - stop_server() diff --git a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs b/litellm-rust/crates/python-bridge/tests/ocr_retained.rs deleted file mode 100644 index 60128ff4a27..00000000000 --- a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs +++ /dev/null @@ -1,51 +0,0 @@ -use pyo3::prelude::*; -use rstest::rstest; -use serial_test::serial; - -#[path = "support/mod.rs"] -mod support; - -use support::native::{native_globals, run_fixture}; - -#[rstest] -#[case::differential_callbacks_wire("differential_callbacks_wire")] -#[case::known_gap_send_time_auth_and_request_hook("known_gap_send_time_auth_and_request_hook")] -#[case::known_gap_custom_transport("known_gap_custom_transport")] -#[case::known_gap_pre_call_timeout_mutation("known_gap_pre_call_timeout_mutation")] -#[case::negative_control_copied_caller_document("negative_control_copied_caller_document")] -#[case::negative_control_rebound_logging_body("negative_control_rebound_logging_body")] -#[case::negative_control_rebound_logging_headers("negative_control_rebound_logging_headers")] -#[case::differential_callback_retained_mutation_after_post_received( - "differential_callback_retained_mutation_after_post_received" -)] -#[case::public_rust_dispatch_wire_fallback_and_escaping_base_exception( - "public_rust_dispatch_wire_fallback_and_escaping_base_exception" -)] -#[case::collection_after_success_and_failures("collection_after_success_and_failures")] -#[case::callback_retained_graph_remains_usable_then_collects( - "callback_retained_graph_remains_usable_then_collects" -)] -#[case::collection_after_cancellation_during_blocked_transport( - "collection_after_cancellation_during_blocked_transport" -)] -#[ignore = "requires repo Python"] -#[serial(python_interpreter)] -fn retained_real_production_boundary_differential_and_lifecycle( - #[case] scenario: &str, -) -> PyResult<()> { - Python::initialize(); - Python::attach(|py| { - let globals = native_globals(py)?; - run_fixture( - py, - &globals, - include_str!("fixtures/ocr_retained.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/ocr_retained.py" - ), - )?; - globals.get_item("run_case")?.unwrap().call1((scenario,))?; - Ok(()) - }) -} diff --git a/litellm-rust/crates/python-bridge/tests/retained_http.rs b/litellm-rust/crates/python-bridge/tests/retained_http.rs deleted file mode 100644 index 2a4ea5c7892..00000000000 --- a/litellm-rust/crates/python-bridge/tests/retained_http.rs +++ /dev/null @@ -1,163 +0,0 @@ -use std::process::Command; -use std::time::{Duration, Instant}; - -use pyo3::prelude::*; -use rstest::rstest; -use serial_test::serial; - -#[path = "support/mod.rs"] -mod support; - -use support::native::{native_globals, run_fixture}; - -#[test] -fn uc_cold_reentry_execution_cache() -> PyResult<()> { - cold_cache_reentry("retained_execution.py", "uc_cold_reentry_execution_cache") -} - -#[test] -fn uc_cold_reentry_callback_cache() -> PyResult<()> { - cold_cache_reentry("retained_callback.py", "uc_cold_reentry_callback_cache") -} - -fn cold_cache_reentry(filename: &str, test: &str) -> PyResult<()> { - if std::env::var("LITELLM_COLD_REENTRY_CHILD").as_deref() != Ok(filename) { - let mut child = Command::new(std::env::current_exe().unwrap()) - .args(["--exact", test, "--nocapture"]) - .env("LITELLM_COLD_REENTRY_CHILD", filename) - .spawn() - .unwrap(); - let deadline = Instant::now() + Duration::from_secs(15); - loop { - if let Some(status) = child.try_wait().unwrap() { - assert!( - status.success(), - "{filename} reentry child failed: {status}" - ); - return Ok(()); - } - if Instant::now() >= deadline { - child.kill().unwrap(); - child.wait().unwrap(); - panic!("{filename} reentry did not complete within 15 seconds"); - } - std::thread::sleep(Duration::from_millis(10)); - } - } - - Python::initialize(); - Python::attach(|py| { - let globals = native_globals(py)?; - run_fixture( - py, - &globals, - include_str!("fixtures/retained_http_contract.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/retained_http_contract.py" - ), - )?; - globals - .get_item("run_cold_cache_reentry")? - .unwrap() - .call1((filename,))?; - Ok(()) - }) -} - -#[test] -#[serial(python_interpreter)] -fn retained_routes_preserve_callbacks_context_wire_and_ownership() -> PyResult<()> { - Python::initialize(); - Python::attach(|py| { - let globals = native_globals(py)?; - run_fixture( - py, - &globals, - include_str!("fixtures/retained_http_contract.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/retained_http_contract.py" - ), - )?; - globals - .get_item("run_ownership_contract")? - .unwrap() - .call0()?; - Ok(()) - }) -} - -#[test] -#[serial(python_interpreter)] -fn retained_routes_surface_transport_failures_as_runtime_error() -> PyResult<()> { - Python::initialize(); - Python::attach(|py| { - let globals = native_globals(py)?; - run_fixture( - py, - &globals, - include_str!("fixtures/retained_http_contract.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/retained_http_contract.py" - ), - )?; - globals.get_item("run_error_contract")?.unwrap().call0()?; - Ok(()) - }) -} - -#[rstest] -#[case::original_retain_mutate("retain_mutate", "original")] -#[case::original_caller_closure("caller_closure", "original")] -#[case::original_delayed_after_prepare_before_encode_read( - "delayed_after_prepare_before_encode_read", - "original" -)] -#[case::original_field_replace("field_replace", "original")] -#[case::original_mutate_then_caught_error_observe("mutate_then_caught_error_observe", "original")] -#[case::copy_caller_inputs_retain_mutate("retain_mutate", "copy_caller_inputs_before_prepare")] -#[case::copy_caller_inputs_caller_closure("caller_closure", "copy_caller_inputs_before_prepare")] -#[case::copy_roots_retain_mutate("retain_mutate", "copy_roots_after_prepare")] -#[case::copy_roots_delayed_after_prepare_before_encode_read( - "delayed_after_prepare_before_encode_read", - "copy_roots_after_prepare" -)] -#[case::copy_roots_mutate_then_caught_error_observe( - "mutate_then_caught_error_observe", - "copy_roots_after_prepare" -)] -#[case::encode_logging_replacements_field_replace("field_replace", "encode_logging_replacements")] -#[case::reconstruct_logging_envelope_retain_mutate("retain_mutate", "reconstruct_logging_envelope")] -#[case::reconstruct_logging_envelope_field_replace("field_replace", "reconstruct_logging_envelope")] -#[case::reconstruct_root_tuple_retain_mutate("retain_mutate", "reconstruct_root_tuple")] -#[case::reconstruct_root_tuple_delayed_after_prepare_before_encode_read( - "delayed_after_prepare_before_encode_read", - "reconstruct_root_tuple" -)] -#[case::reconstruct_root_tuple_field_replace("field_replace", "reconstruct_root_tuple")] -#[serial(python_interpreter)] -fn retained_production_transport_comparison_matrix( - #[case] scenario: &str, - #[case] variant: &str, -) -> PyResult<()> { - Python::initialize(); - Python::attach(|py| { - let globals = native_globals(py)?; - run_fixture( - py, - &globals, - include_str!("fixtures/retained_http_contract.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/retained_http_contract.py" - ), - )?; - globals - .get_item("run_comparison_contract")? - .unwrap() - .call1((scenario, variant))?; - Ok(()) - }) -} diff --git a/litellm-rust/crates/python-bridge/tests/support/mod.rs b/litellm-rust/crates/python-bridge/tests/support/mod.rs deleted file mode 100644 index ce12f2f4f13..00000000000 --- a/litellm-rust/crates/python-bridge/tests/support/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod native; diff --git a/litellm-rust/crates/python-bridge/tests/support/native.rs b/litellm-rust/crates/python-bridge/tests/support/native.rs deleted file mode 100644 index 0023f4fe1c4..00000000000 --- a/litellm-rust/crates/python-bridge/tests/support/native.rs +++ /dev/null @@ -1,21 +0,0 @@ -use pyo3::prelude::*; -use pyo3::types::PyDict; - -pub fn native_globals(py: Python<'_>) -> PyResult> { - let module = pyo3::wrap_pymodule!(_native::_native)(py).into_bound(py); - let globals = PyDict::new(py); - globals.set_item("native", &module)?; - Ok(globals) -} - -pub fn run_fixture( - py: Python<'_>, - globals: &Bound<'_, PyDict>, - source: &str, - filename: &str, -) -> PyResult<()> { - let builtins = py.import("builtins")?; - let code = builtins.call_method1("compile", (source, filename, "exec"))?; - builtins.call_method1("exec", (code, globals))?; - Ok(()) -} diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index e4628fd784e..7a39ddbd607 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs @@ -1,3 +1,6 @@ +use std::process::Command; +use std::time::{Duration, Instant}; + use pyo3::prelude::*; use pyo3::types::PyDict; use rstest::{fixture, rstest}; @@ -11,6 +14,45 @@ mod support; use support::python::{InitializedPython, initialized_python, run_fixture}; +#[test] +fn cold_awaited_adapter_initialization_allows_reentry() -> PyResult<()> { + let test = "cold_awaited_adapter_initialization_allows_reentry"; + let child_env = "LITELLM_INTEROP_COLD_REENTRY_CHILD"; + if std::env::var(child_env).as_deref() != Ok(test) { + let mut child = Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test, "--nocapture"]) + .env(child_env, test) + .spawn() + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(15); + loop { + if let Some(status) = child.try_wait().unwrap() { + assert!( + status.success(), + "awaited adapter reentry child failed: {status}" + ); + return Ok(()); + } + if Instant::now() >= deadline { + child.kill().unwrap(); + child.wait().unwrap(); + panic!("awaited adapter reentry did not complete within 15 seconds"); + } + std::thread::sleep(Duration::from_millis(10)); + } + } + + let globals = scenario_scope(&initialized_python()); + Python::attach(|py| { + globals + .bind(py) + .get_item("cold_awaited_adapter_reentry")? + .unwrap() + .call1((globals.bind(py).get_item("factory")?.unwrap(),))?; + Ok(()) + }) +} + #[fixture] fn scenario_scope(initialized_python: &InitializedPython) -> Py { let _ = initialized_python; 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 ed5201d37d7..d4ff68b343e 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py @@ -2,7 +2,9 @@ import asyncio import contextvars import copy import gc +import inspect import json +import sys import threading import weakref from unittest import TestCase @@ -18,6 +20,65 @@ class Value: pass +def cold_awaited_adapter_reentry(factory): + events = [] + compilations = [] + + def invoke_failure(): + error = LookupError("cold-cache callback error") + + async def failing(): + raise error + + owner = factory.prepare(failing, (), awaited=True) + pending = owner.invoke() + try: + assert inspect.getcoroutinestate(pending) == inspect.CORO_CREATED + with TestCase().assertRaises(LookupError) as caught: + pending.send(None) + assert caught.exception is error + finally: + pending.close() + owner.close() + error.__traceback__ = None + + def audit(event, args): + if event != "compile" or args[1] != "retained_callback.py": + return + compilations.append(args[1]) + if len(compilations) == 1: + events.append("entered") + invoke_failure() + events.append("nested completed") + + sys.addaudithook(audit) + invoke_failure() + events.append("outer completed") + assert events == ["entered", "nested completed", "outer completed"] + assert len(compilations) == 2 + assert factory.live == 0 + + result = object() + direct = factory.prepare(lambda value: value, (result,)) + try: + assert direct.invoke() is result + finally: + direct.close() + + async def successful(value, *, alias): + assert value is alias + await checkpoint() + return value + + owner = factory.prepare(successful, (result,), {"alias": result}, awaited=True) + try: + assert asyncio.run(owner.invoke()) is result + finally: + owner.close() + assert len(compilations) == 2 + assert factory.live == 0 + + class ReferenceFactory: def __init__(self): self.live = 0 diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 2eeef48cb8d..74f192b0115 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -31,7 +31,6 @@ 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.rust_bridge import ocr_retained as rust_ocr_retained_bridge from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -308,44 +307,10 @@ def _map_rust_ocr_error( status_code=status or 500, headers={}, # mutable-ok: provider error factories require a concrete header dict ) - - -def _retained_ocr_boundary(prepared: _PreparedOCRRequest) -> rust_ocr_retained_bridge.OCRRetainedBoundary: - return rust_ocr_retained_bridge.OCRRetainedBoundary( - handler=base_llm_http_handler, - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - custom_llm_provider=prepared.custom_llm_provider, - timeout=prepared.effective_timeout, - ) - - def _run_rust_ocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], - *, - load_retained: Callable[[], rust_ocr_retained_bridge.RustOCRRetained | None] = ( - rust_ocr_retained_bridge.load_rust_ocr_retained - ), ) -> OCRResponse | None: - if prepared_request.custom_llm_provider == "mistral" and not rust_ocr_bridge.has_rust_ocr_override(): - retained: Final = load_retained() - if ( - retained is None - or rust_ocr_retained_bridge.retained_timeout_seconds(prepared_request.effective_timeout) is None - ): - return None - retained_response: Final = retained(_retained_ocr_boundary(prepared_request)) - if retained_response is None: - raise ValueError("Retained OCR returned no response after preparation") - return retained_response if rust_ocr_bridge.load_rust_ocr() is None: return None prepared: Final = _prepare_rust_ocr_call( @@ -373,24 +338,7 @@ def _run_rust_ocr( async def _run_rust_aocr( prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], - *, - load_retained: Callable[[], rust_ocr_retained_bridge.RustAOCRRetained | None] = ( - rust_ocr_retained_bridge.load_rust_aocr_retained - ), ) -> OCRResponse | None: - if prepared_request.custom_llm_provider == "mistral" and not rust_ocr_bridge.has_rust_ocr_override( - asynchronous=True - ): - retained: Final = load_retained() - if ( - retained is None - or rust_ocr_retained_bridge.retained_timeout_seconds(prepared_request.effective_timeout) is None - ): - return None - retained_response: Final = await retained(_retained_ocr_boundary(prepared_request)) - if retained_response is None: - raise ValueError("Retained OCR returned no response after preparation") - return retained_response if rust_ocr_bridge.load_rust_aocr() is None: return None prepared: Final = _prepare_rust_ocr_call( diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 05c09528df6..b7fdb5a98ef 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -53,10 +53,6 @@ _OCR: Final = NativeBinding("ocr", validate=_as_ocr) _AOCR: Final = NativeBinding("aocr", validate=_as_aocr) -def has_rust_ocr_override(*, asynchronous: bool = False) -> bool: - return _rust_aocr_impl is not None if asynchronous else _rust_ocr_impl is not None - - def load_rust_ocr() -> RustOcr | None: return _OCR.load() diff --git a/litellm/rust_bridge/ocr_retained.py b/litellm/rust_bridge/ocr_retained.py deleted file mode 100644 index cc431aabade..00000000000 --- a/litellm/rust_bridge/ocr_retained.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Python preparation, encoding, and response transforms for retained OCR calls.""" - -from __future__ import annotations - -import math -from collections.abc import Awaitable, Callable -from dataclasses import dataclass, field -from typing import Final, cast # noqa: TID251 # native callables and legacy header types require boundary casts - -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.timeouts import timeout_to_seconds - -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 retained_timeout_seconds(timeout: float | httpx.Timeout) -> float | None: - seconds: Final = timeout_to_seconds(timeout) - return seconds if seconds is not None and math.isfinite(seconds) and seconds > 0 else None - - -@dataclass(kw_only=True, slots=True) -class OCRRetainedBoundary: - 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 = retained_timeout_seconds(self.timeout) - if seconds is None: - raise ValueError("Retained OCR requires a positive finite read timeout") - if self.client is None: - raise RuntimeError("Retained OCR 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("Retained OCR 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, - ) - - -RustOCRRetained = Callable[[OCRRetainedBoundary], OCRResponse] -RustAOCRRetained = Callable[[OCRRetainedBoundary], Awaitable[OCRResponse]] - - -def load_rust_ocr_retained() -> RustOCRRetained | None: - from litellm.rust_bridge import get_native_bridge - - native: Final = get_native_bridge() - return cast(RustOCRRetained | None, getattr(native, "ocr_retained", None)) - - -def load_rust_aocr_retained() -> RustAOCRRetained | None: - from litellm.rust_bridge import get_native_bridge - - native: Final = get_native_bridge() - return cast(RustAOCRRetained | None, getattr(native, "aocr_retained", None)) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index bd3a8d186c9..36a29801e1a 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -816,14 +816,23 @@ def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) -@pytest.mark.parametrize("rust_enabled", [False, True], ids=["disabled", "unavailable"]) -def test_ocr_uses_python_when_native_execution_is_unavailable(monkeypatch, rust_enabled): - def load_native(): - assert rust_enabled, "disabled OCR must not load the native module" - return None +def test_ocr_does_not_route_to_rust_when_disabled(): + """With the flag off, the bridge must not be consulted even if an impl exists.""" + bridge = RecordingBridge() + litellm.rust(False) + rust_bridge.set_rust_ocr(ocr=bridge) - monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", load_native) - litellm.rust(rust_enabled) + assert rust_bridge.rust_ocr_enabled() is False + # The impl stays available for injection, but the disabled flag gates usage, + # so ocr() never reaches the Rust path (asserted via the enabled-path test). + assert bridge.calls == [] + + +def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): + """Rust enabled but no bridge available (no injected impl, no compiled wheel): + ocr() must degrade to the Python HTTP handler instead of raising.""" + monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) + litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} 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 a6e330afa26..a7f50a82a99 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -10,18 +10,14 @@ import sys import tempfile import threading import zipfile -from collections import Counter from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from queue import SimpleQueue from socket import socket as Socket -from types import FrameType from typing import Final REQUEST_STARTED: Final = threading.Event() REQUEST_CANCELLED: Final = threading.Event() -PUBLIC_OCR_REQUESTS: Final[SimpleQueue[str]] = SimpleQueue() ANTHROPIC_RESPONSE: Final = ( b'{"id":"msg_native","type":"message","role":"assistant",' @@ -36,24 +32,10 @@ class NativeRouteHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: content_length: Final = int(self.headers.get("content-length", "0")) - wire_body: Final = self.rfile.read(content_length) - body: Final = json.loads(wire_body) + body: Final = json.loads(self.rfile.read(content_length)) route: Final = self.headers.get("x-test-route") outcome: Final = self.headers.get("x-test-outcome") - public_case: Final = self.headers.get("x-test-public-case") - if public_case is not None: - assert self.headers.get("x-test-callback") == public_case - assert self.headers.get("x-test-view-only") is None - assert self.headers.get_all("x-test-callback") == [public_case] - assert wire_body == ( - b'{"model":"mistral-ocr-latest","document":{"type":"document_url",' - b'"document_url":"https://example.com/document.pdf"},"include_image_base64":true,' - b'"document_alias":{"type":"document_url","document_url":"https://example.com/document.pdf"},' - b'"callback_mutation":"observed"}' - ), wire_body assert_native_request(route, outcome, self.path, self.headers, body) - if public_case is not None: - PUBLIC_OCR_REQUESTS.put(public_case) if outcome == "hang": REQUEST_STARTED.set() self.connection.settimeout(5) @@ -245,7 +227,12 @@ 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: @@ -262,149 +249,6 @@ def exercise_routes(native_path: Path, api_base: str) -> object: return native -def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: - import litellm - from litellm.integrations.custom_logger import CustomLogger - from litellm.llms.base_llm.ocr.transformation import OCRResponse - from litellm.rust_bridge import get_native_bridge - from litellm.rust_bridge.ocr_retained import OCRRetainedBoundary - - assert case in {"ocr", "aocr"} - assert Path(litellm.__file__).resolve().is_relative_to(install_root.resolve()) - native: Final = get_native_bridge() - assert native is not None - assert Path(native.__file__).resolve().is_relative_to(install_root.resolve()) - retained: Final = getattr(native, f"{case}_retained") - observed: Final[Counter[str]] = Counter() - roots: Final[dict[str, dict[str, object]]] = {} - phases: Final[list[str]] = [] - document: Final = {"type": "document_url", "document_url": "https://example.com/before-callback.pdf"} - replacement_body: Final = {"replacement": True} - replacement_headers: Final = {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"} - - def observe(frame: FrameType, event: str, arg: object) -> None: - if event == "c_call" and arg is retained: - observed["retained"] += 1 - if event == "call" and frame.f_code is OCRRetainedBoundary.encode.__code__: - observed["encode"] += 1 - - class RetainedLogger(CustomLogger): - def __init__(self, phase: str) -> None: - super().__init__() - self.phase = phase - self.calls = 0 - - def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> dict[str, object]: - self.calls += 1 - additional_args: Final = kwargs["additional_args"] - assert isinstance(additional_args, dict) - if self.phase == "retain": - assert phases == [] - headers: Final = additional_args["headers"] - body: Final = additional_args["complete_input_dict"] - assert isinstance(headers, dict) and isinstance(body, dict) - assert body["model"] == "mistral-ocr-latest" - assert body["include_image_base64"] is False - assert body["document"] is document - assert headers["x-test-callback"] == "before-callback" - roots.update(body=body, headers=headers, view=additional_args) - body["document_alias"] = document - additional_args["complete_input_dict"] = replacement_body - additional_args["headers"] = replacement_headers - else: - assert additional_args is roots["view"] - assert additional_args["complete_input_dict"] is replacement_body - assert additional_args["headers"] is replacement_headers - assert roots["body"]["document_alias"] is document - assert roots["body"]["document"] is document - if self.phase == "mutate": - assert phases == ["retain"] - roots["headers"]["x-test-callback"] = case - roots["body"]["include_image_base64"] = True - document["document_url"] = "https://example.com/document.pdf" - roots["body"]["callback_mutation"] = "observed" - else: - assert self.phase == "observe" - assert phases == ["retain", "mutate"] - assert roots["headers"]["x-test-callback"] == case - assert roots["body"]["include_image_base64"] is True - assert document["document_url"] == "https://example.com/document.pdf" - assert roots["body"]["callback_mutation"] == "observed" - assert replacement_body == {"replacement": True} - assert replacement_headers == { - "x-test-callback": "must-not-send", - "x-test-view-only": "not-on-wire", - } - phases.append(self.phase) - return {"headers": {"x-test-callback": "ignored-return"}, "complete_input_dict": {"invalid": object()}} - - callbacks: Final = tuple(RetainedLogger(phase) for phase in ("retain", "mutate", "observe")) - kwargs: Final = { - "model": "mistral/mistral-ocr-latest", - "document": document, - "api_base": api_base, - "api_key": "sk-native", - "extra_headers": { - "x-test-route": "ocr", - "x-test-outcome": "success", - "x-test-public-case": case, - "x-test-callback": "before-callback", - }, - "include_image_base64": False, - "callbacks": list(callbacks), - "rust": True, - "timeout": 3.0, - "num_retries": 0, - } - previous_profile: Final = sys.getprofile() - sys.setprofile(observe) - try: - response: Final = asyncio.run(litellm.aocr(**kwargs)) if case == "aocr" else litellm.ocr(**kwargs) - finally: - sys.setprofile(previous_profile) - assert isinstance(response, OCRResponse) - assert_success("ocr", response.model_dump()) - assert response.model == "mistral-ocr-latest" - assert tuple(callback.calls for callback in callbacks) == (1, 1, 1) - assert phases == ["retain", "mutate", "observe"], phases - assert observed == {"retained": 1, "encode": 1}, observed - assert roots["view"]["complete_input_dict"] is replacement_body - assert roots["view"]["headers"] is replacement_headers - assert replacement_body == {"replacement": True} - assert replacement_headers == {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"} - assert roots["body"] == { - "model": "mistral-ocr-latest", - "document": document, - "include_image_base64": True, - "document_alias": document, - "callback_mutation": "observed", - } - document["document_url"] = "https://example.com/after-return.pdf" - roots["headers"]["x-after-return"] = "usable" - assert roots["body"]["document"] is roots["body"]["document_alias"] is document - assert roots["headers"]["x-after-return"] == "usable" - assert replacement_body == {"replacement": True} - assert replacement_headers == {"x-test-callback": "must-not-send", "x-test-view-only": "not-on-wire"} - return 0 - - -def verify_public_ocr(wheel: Path, wheel_root: Path, api_base: str) -> None: - install_root: Final = wheel_root / "sdk-venv" - python: Final = install_root / "bin" / "python" - subprocess.run(("uv", "venv", "--python", sys.executable, str(install_root)), check=True) - subprocess.run(("uv", "pip", "install", "--python", str(python), str(wheel.resolve())), check=True) - for case in ("ocr", "aocr"): - subprocess.run( - (str(python), "-I", str(Path(__file__).resolve()), "public-ocr", str(install_root), api_base, case), - cwd=install_root, - env=os.environ | {"LITELLM_LOCAL_MODEL_COST_MAP": "True", "NO_PROXY": "127.0.0.1", "no_proxy": "127.0.0.1"}, - check=True, - timeout=60, - ) - assert PUBLIC_OCR_REQUESTS.get_nowait() == case - assert PUBLIC_OCR_REQUESTS.empty(), f"{case} sent more than one upstream request" - - def exercise_signal(native: object, api_base: str) -> int: try: native.messages( @@ -476,7 +320,6 @@ def verify_wheel(wheel: Path) -> int: api_base: Final = f"http://127.0.0.1:{server.server_address[1]}" try: verify_sigint(native_path, api_base) - verify_public_ocr(wheel, wheel_root, api_base) finally: server.shutdown() server.server_close() @@ -490,8 +333,6 @@ def main() -> int: if len(sys.argv) == 4 and sys.argv[1] == "child": native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3]) return exercise_signal(native, sys.argv[3]) - if len(sys.argv) == 5 and sys.argv[1] == "public-ocr": - return exercise_public_ocr(Path(sys.argv[2]), sys.argv[3], sys.argv[4]) sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") return 2