diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 188d5b8400d..45558be9a8c 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -102,7 +102,7 @@ jobs: restore-keys: | ${{ runner.os }}-cargo-${{ github.job }}- - - run: cargo test --workspace --locked + - run: cargo test --workspace --locked --exclude litellm-python-interop --exclude litellm-python-bridge working-directory: litellm-rust - run: cargo test -p litellm-core --features bedrock-auth --locked diff --git a/Makefile b/Makefile index f5f2193ba86..c70cb026e8f 100644 --- a/Makefile +++ b/Makefile @@ -55,12 +55,9 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" -<<<<<<< HEAD @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" -======= - @echo " make test-rust-python - Run ignored Python-integrated Cargo tests (optional TEST_FILTER=substring)" + @echo " make test-rust-python - Run ignored Python-integrated Cargo tests" @echo " make lint-rust-python-fixtures - Check Rust test Python fixtures with Ruff" ->>>>>>> 1f18773bf2 (wip) @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." @@ -317,7 +314,7 @@ 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 -- --ignored $(if $(TEST_FILTER),"$(TEST_FILTER)") + -p litellm-python-interop -p litellm-python-bridge --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 diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index dfacf37b6cd..3b1e84dd91c 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -187,3 +187,11 @@ cargo test -p litellm-ai-gateway --features server When a Rust path is exposed through Python, add Python parity tests that compare the existing Python output with the Rust-backed output. + +For Python-integrated tests that need the repository's Python environment, run +from the repository root: + +```bash +make test-rust-python +make lint-rust-python-fixtures +``` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 18a08bac8bc..54757ae634f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1474,6 +1474,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "serial_test", "tokio", "tokio-tungstenite", "tracing", @@ -1488,6 +1489,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "serial_test", ] [[package]] @@ -1496,6 +1498,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1605,6 +1616,29 @@ dependencies = [ "winapi", ] +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1933,6 +1967,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "regex" version = "1.13.1" @@ -2177,6 +2220,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "sct" version = "0.7.1" @@ -2282,6 +2331,28 @@ dependencies = [ "serde", ] +[[package]] +name = "serial_test" +version = "4.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6df5ed973ad8d834e09f824f9e9f449af6b9a3745f78dec7cc752770bd3bf11" +dependencies = [ + "once_cell", + "parking_lot", + "serial_test_derive", +] + +[[package]] +name = "serial_test_derive" +version = "4.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a22144e767da4ddd8416dbf383700542ffd8a5dc493dfecedfe1fe3ad03c98ae" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.0", +] + [[package]] name = "sha1" version = "0.10.7" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 41c2128b500..441ad33583f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -30,10 +30,10 @@ reqwest = { version = "0.12", default-features = false, features = ["blocking", rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } rustls-native-certs = "0.8" +serial_test = { version = "4.0.1", default-features = false } serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } sha2 = "0.10" -strum = { version = "0.26", features = ["derive"] } subtle = "2" thiserror = "2.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 24a5da7aa9c..a5629744574 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -57,47 +57,26 @@ runs for changes under `litellm-rust/`. ### Python-Integrated Tests -From the repository root, use the same entrypoint as the Rust CI workflow to run -the ignored Cargo tests that need the repository's Python dependencies: +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: ```bash make test-rust-python -make test-rust-python TEST_FILTER=component_contract -make test-rust-python TEST_FILTER=retained make lint-rust-python-fixtures ``` -The test target covers `litellm-python-interop` (including `component_contract` -and `prepared_call`) and `litellm-python-bridge` (including `ocr_retained`). -`TEST_FILTER` is an optional Rust test-name substring, not a Python fixture or -Cargo test-binary name. An unmatched filter runs zero tests, so check the test -counts. The underlying command is: +`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. -```bash -cargo test --manifest-path litellm-rust/Cargo.toml \ - -p litellm-python-interop -p litellm-python-bridge --tests --locked -- --ignored -``` - -Install uv and the repository's pinned Rust toolchain first. The -`install-rust-python-test-deps` prerequisite runs -`uv sync --inexact --frozen --no-default-groups --no-install-project` on each -invocation. This installs the locked SDK dependencies without building or -installing LiteLLM, pulling in the full dev groups, or pruning existing venv -packages. No wheel is needed for these embedded-Python tests; the existing wheel -lane checks the installed public interface separately - -The target gets the project interpreter from `uv run --no-sync python`, sets -`PYO3_PYTHON` to that executable, and queries its `sysconfig` for both Python and -platform-specific site-packages. It prepends the repository root and those paths -to `PYTHONPATH`, preserving any existing entries, so embedded Python imports the -checkout and its dependencies. Use uv's `UV_PYTHON` and `UV_PROJECT_ENVIRONMENT` -settings to select a different interpreter or project venv. The target also sets -`LITELLM_LOCAL_MODEL_COST_MAP=True` to use the checked-in model cost map - -Fixture checks are separate from the test target so filtered reruns stay focused. -`make lint-rust-python-fixtures` runs pinned Ruff lint and format checks with -`ruff-tests.toml` over both crates' `tests/` directories, including -`crates/python-bridge/tests/fixtures/ocr_retained.py`, without syncing the project -environment. CI runs both targets; its Python path trigger covers `litellm/**` -so changes to OCR, bridge, logging, streaming, and their shared imports rerun the -integrated tests +The callback lifecycle and retained OCR 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 +exclusive scenario. Unannotated tests do not participate in this isolation. +Keep the attribute below `#[rstest]` so generated cases acquire it before +fixture setup and Python attachment. Tasks and threads inside each scenario +still run concurrently. Separate test processes have separate interpreters, +so these attributes need no cross-process lock when using nextest diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index b54df02d72b..ec1a2988e75 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -32,7 +32,7 @@ pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; -pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; +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/ocr/transport.rs b/litellm-rust/crates/core/src/http_utils/buffered_post.rs similarity index 73% rename from litellm-rust/crates/core/src/ocr/transport.rs rename to litellm-rust/crates/core/src/http_utils/buffered_post.rs index 9501645cb1d..3272beaa7e6 100644 --- a/litellm-rust/crates/core/src/ocr/transport.rs +++ b/litellm-rust/crates/core/src/http_utils/buffered_post.rs @@ -3,7 +3,7 @@ use std::time::Duration; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use crate::constants::OCR_CONNECT_TIMEOUT_SECS; +use crate::constants::BUFFERED_POST_CONNECT_TIMEOUT_SECS; use crate::error::Error; pub struct Request { @@ -23,13 +23,13 @@ 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("OCR timeout must be positive and finite".into()))?; + .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 OCR header name".into()))?; + .map_err(|_| Error::InvalidRequest("invalid header name".into()))?; let value = HeaderValue::from_bytes(&value) - .map_err(|_| Error::InvalidRequest("invalid OCR header value".into()))?; + .map_err(|_| Error::InvalidRequest("invalid header value".into()))?; headers.append(name, value); } @@ -37,7 +37,7 @@ pub async fn send(request: Request) -> Result { let client = CLIENT .get_or_init(|| { reqwest::Client::builder() - .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) + .connect_timeout(Duration::from_secs(BUFFERED_POST_CONNECT_TIMEOUT_SECS)) .redirect(reqwest::redirect::Policy::none()) .no_gzip() .no_brotli() @@ -46,7 +46,7 @@ pub async fn send(request: Request) -> Result { .build() }) .as_ref() - .map_err(|_| Error::Network("could not initialize OCR HTTP client".into()))?; + .map_err(|_| Error::Network("could not initialize HTTP client".into()))?; let response = client .post(request.url) .headers(headers) @@ -54,7 +54,7 @@ pub async fn send(request: Request) -> Result { .timeout(timeout) .send() .await - .map_err(|_| Error::Network("OCR transport failed".into()))?; + .map_err(|_| Error::Network("transport failed".into()))?; let status = response.status().as_u16(); let headers = response .headers() @@ -64,7 +64,7 @@ pub async fn send(request: Request) -> Result { let content = response .bytes() .await - .map_err(|_| Error::Network("could not read OCR response".into()))? + .map_err(|_| Error::Network("could not read response".into()))? .to_vec(); Ok(Response { status, diff --git a/litellm-rust/crates/core/src/ocr/transport/tests.rs b/litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs similarity index 86% rename from litellm-rust/crates/core/src/ocr/transport/tests.rs rename to litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs index 4a77700e9f9..5237e91b66e 100644 --- a/litellm-rust/crates/core/src/ocr/transport/tests.rs +++ b/litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs @@ -18,17 +18,17 @@ async fn rejects_invalid_timeouts_and_headers_without_echoing_wire_values() { }) .await; assert!( - matches!(result, Err(Error::InvalidRequest(message)) if message == "OCR timeout must be positive and finite") + 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 OCR header name", + "invalid header name", ), ( vec![(b"x-proof".to_vec(), b"private\nvalue".to_vec())], - "invalid OCR header value", + "invalid header value", ), ] { let result = send(Request { @@ -39,6 +39,6 @@ async fn rejects_invalid_timeouts_and_headers_without_echoing_wire_values() { assert!(matches!(result, Err(Error::InvalidRequest(message)) if message == expected)); } assert!( - matches!(send(request()).await, Err(Error::Network(message)) if message == "OCR transport failed") + matches!(send(request()).await, Err(Error::Network(message)) if message == "transport failed") ); } diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils/mod.rs similarity index 99% rename from litellm-rust/crates/core/src/http_utils.rs rename to litellm-rust/crates/core/src/http_utils/mod.rs index cb472dd5a57..93dfc1a5a8f 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils/mod.rs @@ -1,5 +1,7 @@ //! Header and upstream-body helpers shared by every route module. +pub mod buffered_post; + use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index c2e707160b6..ec2fbb969a6 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,3 +1,2 @@ pub mod transformation; -pub mod transport; pub mod types; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 42282ca4da4..ccf8e8bd7a3 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,6 @@ 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 5959ecd73c1..593f6d3cfca 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -35,6 +35,7 @@ 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 5b423a32e7d..c3da372ed14 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -89,12 +89,19 @@ where F: Future> + Send + 'static, { pyo3_async_runtimes::tokio::future_into_py(py, async move { - let result = catch_future_panic(future).await?; - let result = map_core_result(result, map_error)?; + let result = run_async_value(future, map_error).await?; 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), diff --git a/litellm-rust/crates/python-bridge/src/routes/bindings.rs b/litellm-rust/crates/python-bridge/src/routes/bindings.rs deleted file mode 100644 index 9c0912d5b2a..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/bindings.rs +++ /dev/null @@ -1,39 +0,0 @@ -use litellm_python_interop::InvocationMode; - -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, - }, - } - } -} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 05b11935956..2873a9a8302 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -7,11 +7,11 @@ mod definition; mod gateway_messages; mod audio_transcription; -mod bindings; mod chat_completions; mod messages; mod ocr; mod ocr_retained; +mod retained_http; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { ocr::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 index ec35d57418b..9b36878013c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs @@ -1,141 +1,15 @@ -use litellm_core::ocr::transport::{self, Request, Response}; -use litellm_python_interop::{InvocationOutcome, PreparedCall}; use pyo3::prelude::*; -use pyo3::sync::PyOnceLock; -use pyo3::types::{PyBytes, PyList, PyTuple}; -use super::bindings::{BoundaryMethod, MethodBinding}; -use crate::errors::core_error_to_pyerr; -use crate::execution::{catch_future_panic, run_sync_value}; - -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()), - ) -} - -fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult { - type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; - let encoded = invoke( - boundary, - BoundaryMethod::Encode.resolve(false), - PyTuple::new(boundary.py(), [roots])?, - )?; - let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) = - encoded.extract(boundary.py())?; - 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>( - boundary: &Bound<'py, PyAny>, - roots: &Bound<'py, PyAny>, -) -> PyResult> { - let request = encode(boundary, roots)?; - pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move { - let response = catch_future_panic(transport::send(request)) - .await? - .map_err(core_error_to_pyerr)?; - 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])?, - ) -} +use super::retained_http; #[pyfunction] fn ocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult> { - let py = boundary.py(); - let roots = prepare(boundary, false)?; - let request = encode(boundary, roots.bind(py))?; - let response = run_sync_value(py, transport::send(request), core_error_to_pyerr)?; - let wire = Wire(response).into_pyobject(py)?; - finish(boundary, wire.as_any(), false) + retained_http::run_sync(boundary) } #[pyfunction] fn aocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult> { - static DRIVER: PyOnceLock> = PyOnceLock::new(); - let py = boundary.py(); - let driver = DRIVER.get_or_try_init(py, || { - PyModule::from_code( - py, - c"async def drive(boundary, prepare, send, finish): - roots = await prepare(boundary, True) - wire = await send(boundary, roots) - return await finish(boundary, wire, True) -", - c"ocr_retained_driver.py", - c"_ocr_retained_driver", - )? - .getattr("drive") - .map(Bound::unbind) - })?; - driver.call1( - py, - ( - boundary, - wrap_pyfunction!(prepare, py)?, - wrap_pyfunction!(send, py)?, - wrap_pyfunction!(finish, py)?, - ), - ) + retained_http::run_async(boundary) } pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/litellm-rust/crates/python-bridge/src/routes/retained_http.rs b/litellm-rust/crates/python-bridge/src/routes/retained_http.rs new file mode 100644 index 00000000000..0ac67a50bb6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/retained_http.rs @@ -0,0 +1,172 @@ +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()), + ) +} + +fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult { + type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; + let encoded = invoke( + boundary, + BoundaryMethod::Encode.resolve(false), + PyTuple::new(boundary.py(), [roots])?, + )?; + let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) = + encoded.extract(boundary.py())?; + 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>( + boundary: &Bound<'py, PyAny>, + roots: &Bound<'py, PyAny>, +) -> PyResult> { + let request = encode(boundary, roots)?; + pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move { + let response = run_async_value(buffered_post::send(request), core_error_to_pyerr).await?; + Ok(Wire(response)) + }) +} + +#[pyfunction] +fn finish( + boundary: &Bound<'_, PyAny>, + wire: &Bound<'_, PyAny>, + asynchronous: bool, +) -> PyResult> { + 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 roots = prepare(boundary, false)?; + let request = encode(boundary, roots.bind(py))?; + let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?; + let wire = Wire(response).into_pyobject(py)?; + finish(boundary, wire.as_any(), false) +} + +pub(crate) fn run_async(boundary: &Bound<'_, PyAny>) -> PyResult> { + static DRIVER: PyOnceLock> = PyOnceLock::new(); + let py = boundary.py(); + let driver = DRIVER.get_or_try_init(py, || { + PyModule::from_code( + py, + c"async def drive(boundary, prepare, send, finish): + roots = await prepare(boundary, True) + wire = await send(boundary, roots) + return await finish(boundary, wire, True) +", + c"retained_http_driver.py", + c"_retained_http_driver", + )? + .getattr("drive") + .map(Bound::unbind) + })?; + driver.call1( + py, + ( + boundary, + wrap_pyfunction!(prepare, py)?, + wrap_pyfunction!(send, py)?, + wrap_pyfunction!(finish, py)?, + ), + ) +} diff --git a/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py b/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py index 6e002970803..6c975c090ab 100644 --- a/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py +++ b/litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py @@ -187,7 +187,7 @@ class RealBoundaryTests(unittest.TestCase): 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} + optional = {"unknown_python_json": {7: ("tuple", 2)}, "nested": nested, "nested_alias": nested["values"]} retained = {} events = [] @@ -202,6 +202,7 @@ class RealBoundaryTests(unittest.TestCase): 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" @@ -218,6 +219,9 @@ class RealBoundaryTests(unittest.TestCase): 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") @@ -225,8 +229,15 @@ class RealBoundaryTests(unittest.TestCase): 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") @@ -240,6 +251,7 @@ class RealBoundaryTests(unittest.TestCase): 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"): @@ -259,16 +271,28 @@ class RealBoundaryTests(unittest.TestCase): 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]},"body_mutation":true,"before_error":true}' + 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() @@ -373,6 +397,7 @@ class RealBoundaryTests(unittest.TestCase): 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() @@ -487,6 +512,8 @@ class RealBoundaryTests(unittest.TestCase): 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( @@ -514,26 +541,29 @@ class RealBoundaryTests(unittest.TestCase): 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 != "encoding")) + 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_encoding_failure_and_http_error(self): + def test_collection_after_success_and_failures(self): async def exercise(): - for outcome in ("success", "encoding", "http"): + 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), "request graph leaked after " + outcome) + 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) @@ -542,21 +572,29 @@ class RealBoundaryTests(unittest.TestCase): def test_callback_retained_graph_remains_usable_then_collects(self): async def exercise(): - for mode in ("native-sync", "native-async"): - with self.subTest(mode=mode): - retained = [] - refs, _ = await self.lifecycle(mode, "success", retained) - gc.collect() - self.assertIsNone(refs[1]()) - self.assertTrue(all(refs[index]() is not None for index in (0, 2, 3, 4))) - retained[0]["document"]["after_return"] = "usable" - retained[0]["sentinel"]["alive"] = "still usable" - retained[1]["X-After-Return"] = "usable" - self.assertEqual(refs[0]()["after_return"], "usable") - self.assertEqual(refs[3]()["alive"], "still usable") - retained.clear() - gc.collect() - self.assertTrue(all(ref() is None for ref in refs)) + 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()) 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 new file mode 100644 index 00000000000..a3399ca5d55 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py @@ -0,0 +1,269 @@ +import asyncio +import contextvars +import gc +import http.server +import inspect +import threading +import weakref + +native = globals()["native"] +started = threading.Event() +release = threading.Event() +server = None +url = None +context = contextvars.ContextVar("proof", default="unset") + + +def start_server(): + global server, 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) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url = "http://127.0.0.1:%s/ocr" % server.server_port + return requests + + +requests = start_server() + + +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 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) + 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: + release.set() + server.shutdown() + server.server_close() + + +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"\x00", self.timeout) + + def finish(self, wire): + raise AssertionError("client-side failure must not reach finish") + + +def run_error_contract(): + cases = [ + ("timeout", "http://10.255.255.1:9/", 0.05), + ("refused", "http://127.0.0.1:1/", 1.0), + ] + 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") diff --git a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs b/litellm-rust/crates/python-bridge/tests/ocr_retained.rs index 6f9430dddf3..c1c7d1e78c9 100644 --- a/litellm-rust/crates/python-bridge/tests/ocr_retained.rs +++ b/litellm-rust/crates/python-bridge/tests/ocr_retained.rs @@ -1,6 +1,11 @@ use pyo3::prelude::*; -use pyo3::types::PyDict; 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")] @@ -13,9 +18,7 @@ use rstest::rstest; #[case::public_rust_dispatch_wire_fallback_and_escaping_base_exception( "public_rust_dispatch_wire_fallback_and_escaping_base_exception" )] -#[case::collection_after_success_encoding_failure_and_http_error( - "collection_after_success_encoding_failure_and_http_error" -)] +#[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" )] @@ -23,257 +26,30 @@ use rstest::rstest; "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 module = pyo3::wrap_pymodule!(_native::_native)(py).into_bound(py); - let globals = PyDict::new(py); - globals.set_item("native", module)?; - let builtins = py.import("builtins")?; - let code = builtins.call_method1( - "compile", - ( - include_str!("fixtures/ocr_retained.py"), - concat!( - env!("CARGO_MANIFEST_DIR"), - "/tests/fixtures/ocr_retained.py" - ), - "exec", + 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" ), )?; - builtins.call_method1("exec", (code, &globals))?; - globals + let case = globals .get_item("RealBoundaryTests")? .unwrap() - .call1((format!("test_{scenario}"),))? - .call_method0("debug")?; + .call1((format!("test_{scenario}"),))?; + let outcome = case.call_method0("debug"); + let cleanup = case.call_method0("doCleanups"); + outcome?; + cleanup?; Ok(()) }) } - -#[test] -fn retained_routes_preserve_callbacks_context_wire_and_ownership() -> PyResult<()> { - Python::initialize(); - Python::attach(|py| { - let module = pyo3::wrap_pymodule!(_native::_native)(py).into_bound(py); - let globals = PyDict::new(py); - globals.set_item("native", module)?; - py.run( - cr" -import asyncio -import contextvars -import gc -import http.server -import inspect -import threading -import weakref - -started = threading.Event() -release = threading.Event() -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) -thread = threading.Thread(target=server.serve_forever, daemon=True) -thread.start() -url = 'http://127.0.0.1:%s/ocr' % server.server_port -context = contextvars.ContextVar('proof', default='unset') - -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 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() - -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) - 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: - release.set() - server.shutdown() - server.server_close() - thread.join(5) -", - Some(&globals), - Some(&globals), - ) - }) -} diff --git a/litellm-rust/crates/python-bridge/tests/retained_http.rs b/litellm-rust/crates/python-bridge/tests/retained_http.rs new file mode 100644 index 00000000000..e7046250ed1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/retained_http.rs @@ -0,0 +1,47 @@ +use pyo3::prelude::*; + +#[path = "support/mod.rs"] +mod support; + +use support::native::{native_globals, run_fixture}; + +#[test] +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] +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(()) + }) +} diff --git a/litellm-rust/crates/python-bridge/tests/support/mod.rs b/litellm-rust/crates/python-bridge/tests/support/mod.rs new file mode 100644 index 00000000000..ce12f2f4f13 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/support/mod.rs @@ -0,0 +1 @@ +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 new file mode 100644 index 00000000000..0023f4fe1c4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/support/native.rs @@ -0,0 +1,21 @@ +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/Cargo.toml b/litellm-rust/crates/python-interop/Cargo.toml index 9da6af6e2e2..3e7d87dc760 100644 --- a/litellm-rust/crates/python-interop/Cargo.toml +++ b/litellm-rust/crates/python-interop/Cargo.toml @@ -12,4 +12,5 @@ serde.workspace = true [dev-dependencies] rstest.workspace = true +serial_test.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index 61cb1b6b44b..37fc2fad15f 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs @@ -1,30 +1,15 @@ use pyo3::prelude::*; use pyo3::types::PyDict; use rstest::{fixture, rstest}; +use serial_test::{parallel, serial}; #[path = "support/callback_owner.rs"] mod callback_owner; -struct InitializedPython; +#[path = "support/mod.rs"] +mod support; -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(()) -} - -#[fixture] -#[once] -fn initialized_python() -> InitializedPython { - Python::initialize(); - InitializedPython -} +use support::python::{InitializedPython, initialized_python, run_fixture}; #[fixture] fn scenario_scope(initialized_python: &InitializedPython) -> Py { @@ -65,6 +50,10 @@ fn scenario_scope(initialized_python: &InitializedPython) -> Py { #[case::stream_lifecycle("stream_lifecycle")] #[case::sync_stream_lifecycle("sync_stream_lifecycle")] #[case::repeated_ownership("repeated_ownership")] +#[case::retained_field_replacement("retained_field_replacement")] +#[case::queued_graph_ownership("queued_graph_ownership")] +#[case::detached_work_after_error("detached_work_after_error")] +#[serial(python_interpreter)] fn lifecycle_contract( scenario_scope: Py, #[case] scenario: &str, @@ -86,10 +75,15 @@ fn lifecycle_contract( #[rstest] #[case::real_async_logging("real_async_logging")] +#[case::real_pre_call_logging("real_pre_call_logging")] +#[case::real_copy_boundaries("real_copy_boundaries")] +#[case::real_logging_worker("real_logging_worker")] +#[case::real_sync_stream_copies("real_sync_stream_copies")] #[case::real_stream_completion("real_stream_completion")] #[case::real_stream_close("real_stream_close")] #[case::real_stream_cancellation("real_stream_cancellation")] #[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] +#[serial(python_interpreter)] fn component_contract( scenario_scope: Py, #[case] scenario: &str, @@ -116,6 +110,40 @@ fn component_contract( } #[rstest] +#[case::real_logging_queue_chain("real_logging_queue_chain")] +#[case::real_crowdstrike_translator_identity("real_crowdstrike_translator_identity")] +#[case::real_rubrik_block_lifecycle("real_rubrik_block_lifecycle")] +#[case::real_parallel_guardrail_snapshots("real_parallel_guardrail_snapshots")] +#[case::real_purview_sync_background("real_purview_sync_background")] +#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] +#[serial(python_interpreter)] +fn integration_contract( + scenario_scope: Py, + #[case] scenario: &str, + #[values(false, true)] retained: bool, +) -> PyResult<()> { + Python::attach(|py| { + let globals = scenario_scope.bind(py); + run_fixture( + py, + globals, + include_str!("fixtures/callback_integrations.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/callback_integrations.py" + ), + )?; + globals.get_item("run_scenario")?.unwrap().call1(( + scenario, + retained, + globals.get_item("factory")?.unwrap(), + ))?; + Ok(()) + }) +} + +#[rstest] +#[parallel(python_interpreter)] fn detached_release(initialized_python: &InitializedPython) -> PyResult<()> { use litellm_python_interop::{InvocationMode, PreparedCall}; use pyo3::types::PyTuple; @@ -148,6 +176,7 @@ fn detached_release(initialized_python: &InitializedPython) -> PyResult<()> { #[rstest] #[case::direct(false)] #[case::awaited(true)] +#[parallel(python_interpreter)] fn outcome_identifies_binding(scenario_scope: Py, #[case] awaited: bool) -> PyResult<()> { use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; use pyo3::types::PyTuple; @@ -178,6 +207,7 @@ fn outcome_identifies_binding(scenario_scope: Py, #[case] awaited: bool) } #[rstest] +#[parallel(python_interpreter)] fn awaited_raise_surfaces_when_driven(scenario_scope: Py) -> PyResult<()> { use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; use pyo3::types::PyTuple; @@ -188,7 +218,7 @@ fn awaited_raise_surfaces_when_driven(scenario_scope: Py) -> PyResult<() Some(globals), None, )?; - let started = || -> PyResult { Ok(globals.get_item("events")?.unwrap().len()?) }; + let started = || -> PyResult { globals.get_item("events")?.unwrap().len() }; let call = PreparedCall::new( InvocationMode::Await, globals.get_item("callback")?.unwrap().unbind(), diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py index 5c42da2b634..6789d962c53 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py @@ -1,14 +1,23 @@ import asyncio +import atexit +import contextvars +import gc +import json +import threading +import weakref from datetime import datetime from unittest import TestCase +import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import safe_deep_copy from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.logging_worker import LoggingWorker from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices, Usage -def logger_for(callbacks=(), stream=False): +def logger_for(callbacks=(), stream=False, input_callbacks=(), sync_callbacks=()): return Logging( model="test", messages=[{"role": "user", "content": "test"}], @@ -18,9 +27,63 @@ def logger_for(callbacks=(), stream=False): litellm_call_id="retained-test", function_id="retained-test", dynamic_async_success_callbacks=list(callbacks), + dynamic_input_callbacks=list(input_callbacks), + dynamic_success_callbacks=list(sync_callbacks), ) +async def real_pre_call_logging(owners): + retained = [] + observed = [] + ignored = {"replacement": True} + metadata = {"secret": "private", "keep": []} + removed = object() + lock = threading.Lock() + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained.append(kwargs) + return ignored + + class Mutate(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["normalized"] = "normalized" + assert kwargs.pop("remove") is removed + kwargs["retained_metadata"]["secret"] = "masked" + return ignored + + class Fail(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["lock"] = lock + kwargs["retained_metadata"]["keep"].append("before failure") + raise RuntimeError("expected pre-call callback failure") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append((kwargs, messages)) + + logger = logger_for(input_callbacks=[Retain(), Mutate(), Fail(), Observe()]) + details = logger.model_call_details + details.update(retained_metadata=metadata, normalized=None, remove=removed) + messages = logger.messages + additional = {"headers": {"test": "header"}} + owner = owners.prepare(logger.pre_call, (messages, "test-key"), {"additional_args": additional}) + try: + assert owner.invoke() is None + finally: + owner.close() + assert retained == [details] and observed == [(details, messages)] + assert retained[0] is details and observed[0][0] is details + assert observed[0][1] is messages and details["input"] is messages + assert details["additional_args"] is additional + assert details["retained_metadata"] is metadata + assert metadata == {"secret": "masked", "keep": ["before failure"]} + assert details["normalized"] == "normalized" and "remove" not in details + assert details["lock"] is lock and "replacement" not in details + with TestCase().assertRaises(TypeError): + json.dumps({"lock": details["lock"]}) + + async def real_async_logging(owners): observations = [] task = asyncio.current_task() @@ -28,6 +91,8 @@ async def real_async_logging(owners): result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}]) replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}]) shared = {} + side_channel = {} + replaced_kwargs = [] class Retain(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): @@ -39,13 +104,17 @@ async def real_async_logging(owners): asyncio.get_running_loop().call_soon(gate.set) await gate.wait() kwargs["retained_shared"]["changed"] = True + side_channel["failed_hook"] = kwargs result.choices[0].message.content = "mutated" raise RuntimeError("expected async callback failure") class Replace(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): observations.append(("replace", kwargs, result)) - return kwargs, replacement + updated = {**kwargs, "adopted": True} + replaced_kwargs.append(updated) + side_channel["replacement"] = replacement + return updated, replacement class Observe(CustomLogger): async def async_logging_hook(self, kwargs, result, call_type): @@ -65,7 +134,125 @@ async def real_async_logging(owners): assert observations[0][2] is result and observations[1][2] is result assert observations[2][2] is replacement and observations[3][2] is replacement assert observations[0][1]["retained_shared"] is shared and shared["changed"] + assert observations[0][1] is observations[1][1] is side_channel["failed_hook"] + assert observations[2][1] is observations[3][1] is logger.model_call_details is replaced_kwargs[0] + assert logger.model_call_details is not observations[0][1] + assert logger.model_call_details["adopted"] and "adopted" not in observations[0][1] + assert logger.model_call_details["retained_shared"] is shared + assert side_channel["replacement"] is replacement assert result.choices[0].message.content == "mutated" + observations[0][1]["retained_shared"]["after_replacement"] = True + assert observations[3][1]["retained_shared"]["after_replacement"] + + +async def real_copy_boundaries(owners): + lock = threading.Lock() + shared = {"values": []} + standard = { + "messages": [{"role": "user", "content": "private"}], + "response": {"choices": [{"message": {"content": "private"}}]}, + "metadata": shared, + } + details = {"standard_logging_object": standard, "shared": shared, "lock": lock} + passthrough = owners.prepare(CustomLogger().redact_standard_logging_payload_from_model_call_details, (details,)) + try: + assert passthrough.invoke() is details + finally: + passthrough.close() + logger = CustomLogger(turn_off_message_logging=True) + redact = owners.prepare(logger.redact_standard_logging_payload_from_model_call_details, (details,)) + try: + redacted = redact.invoke() + finally: + redact.close() + assert redacted is not details + assert redacted["standard_logging_object"] is not standard + assert redacted["shared"] is shared and redacted["lock"] is lock + assert redacted["standard_logging_object"]["metadata"] is shared + redacted["standard_logging_object"]["metadata"]["values"].append("shared mutation") + assert shared["values"] == ["shared mutation"] + assert redacted["standard_logging_object"]["messages"][0]["content"] == "redacted-by-litellm" + assert redacted["standard_logging_object"]["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert standard["messages"][0]["content"] == "private" + assert standard["response"]["choices"][0]["message"]["content"] == "private" + + original_mode = litellm.safe_memory_mode + try: + for safe_mode in (False, True): + litellm.safe_memory_mode = safe_mode + uncopyable = {"lock": lock, "values": []} + values = [] + data = {"copyable": {"values": values, "alias": values}, "uncopyable": uncopyable} + owner = owners.prepare(safe_deep_copy, (data,)) + try: + copied = owner.invoke() + finally: + owner.close() + assert (copied is data) is safe_mode + assert (copied["copyable"] is data["copyable"]) is safe_mode + assert copied["copyable"]["values"] is copied["copyable"]["alias"] + assert (copied["copyable"]["values"] is values) is safe_mode + assert copied["uncopyable"] is uncopyable and copied["uncopyable"]["lock"] is lock + copied["copyable"]["values"].append("copy") + assert copied["copyable"]["alias"] == ["copy"] + copied["uncopyable"]["values"].append("fallback") + assert data["copyable"]["values"] == (["copy"] if safe_mode else []) + assert uncopyable["values"] == ["fallback"] + finally: + litellm.safe_memory_mode = original_mode + + +async def real_logging_worker(owners): + context = contextvars.ContextVar("component_worker_context", default="outside") + entered, release = asyncio.Event(), asyncio.Event() + observations = [] + worker = LoggingWorker(timeout=5, concurrency=1) + + class Payload: + pass + + async def upload(value, *, alias): + assert value is alias + assert context.get() == "submitted" + entered.set() + await release.wait() + observations.append((value.changed, context.get())) + context.set("worker only") + + payload = Payload() + payload.changed = False + reference = weakref.ref(payload) + invocation = owners.prepare(upload, (payload,), {"alias": payload}, awaited=True) + pending = invocation.invoke() + invocation.close() + enqueue = owners.prepare(worker.ensure_initialized_and_enqueue, (pending,)) + stop = owners.prepare(worker.stop, (), awaited=True) + flush = owners.prepare(worker.flush, (), awaited=True) + token = context.set("submitted") + try: + enqueue.invoke() + enqueue.close() + del pending, payload + context.set("consumer") + await entered.wait() + assert reference() is not None + reference().changed = True + release.set() + await flush.invoke() + assert observations == [(True, "submitted")] + assert context.get() == "consumer" + finally: + release.set() + enqueue.close() + flush.close() + await stop.invoke() + stop.close() + context.reset(token) + atexit.unregister(worker._flush_on_exit) + assert worker._worker_task is None and not worker._running_tasks and not worker._dequeued_tasks + assert worker._queue.empty() + gc.collect() + assert reference() is None class ControlledStream: @@ -95,20 +282,16 @@ class ControlledStream: self.closed += 1 -async def drain_component_tasks(): - pending = asyncio.all_tasks() - {asyncio.current_task()} - if pending: - await asyncio.gather(*pending) - - async def real_stream_completion(owners): logger = logger_for(stream=True) completions = [] cached = [] + cache_done = asyncio.Event() class CacheRecorder: async def _add_streaming_response_to_cache(self, response): cached.append(response) + cache_done.set() logger._llm_caching_handler = CacheRecorder() @@ -117,6 +300,7 @@ async def real_stream_completion(owners): logger._on_deferred_stream_complete = complete stream = ControlledStream() + stream.originals[1].usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) wrapper = CustomStreamWrapper( completion_stream=stream, model="test", logging_obj=logger, custom_llm_provider="bedrock" ) @@ -128,7 +312,14 @@ async def real_stream_completion(owners): try: chunk = await pull.invoke() chunks.append(chunk) + if len(chunks) == 1: + assert wrapper.chunks[-1] is chunk + chunk.choices[0].delta.content = "retained hello" if chunk.choices and chunk.choices[0].finish_reason: + stored = wrapper.chunks[-1] + assert stored is not chunk and stored is not stream.originals[1] + assert stored.usage is stream.originals[1].usage + assert getattr(chunk, "usage", None) is None and stored.usage.total_tokens == 2 retained_hidden = chunk._hidden_params hidden_owner = owners.prepare(lambda value: value, (retained_hidden,)) assert not completions @@ -140,12 +331,22 @@ async def real_stream_completion(owners): assert completions == [] response, cache_hit = logger._deferred_stream_complete_args assert response.usage.total_tokens == 8 + assert response.choices[0].message.content == "retained hello" + assert retained_hidden["usage"] is response.usage + usage_chunk = wrapper.chunks[-1] + assert usage_chunk is not stream.originals[-1] + assert usage_chunk.usage is stream.originals[-1].usage + stream.originals[-1].usage.total_tokens = 13 + assert usage_chunk.usage.total_tokens == 13 + assert response.usage.total_tokens == 8 deferred = owners.prepare(logger._on_deferred_stream_complete, (response, cache_hit), awaited=True) logger._on_deferred_stream_complete = None logger._deferred_stream_complete_args = None close = owners.prepare(wrapper.aclose, (), awaited=True) await close.invoke() + await close.invoke() close.close() + assert stream.closed == 1 del wrapper, logger await deferred.invoke() deferred.close() @@ -153,11 +354,95 @@ async def real_stream_completion(owners): assert retained_hidden["usage"].total_tokens == 8 assert hidden_owner.invoke() is retained_hidden hidden_owner.close() - await drain_component_tasks() + await cache_done.wait() assert len(cached) == 1 and cached[0] is not response assert cached[0].choices[0] is not response.choices[0] cached[0].choices[0].message.content = "cache only" - assert response.choices[0].message.content == "hello" + assert response.choices[0].message.content == "retained hello" + + +async def real_sync_stream_copies(owners): + original_disable = litellm.disable_streaming_logging + copy_attempts = [] + + class Uncopyable: + def __deepcopy__(self, memo): + copy_attempts.append(True) + raise RuntimeError("expected streaming deepcopy failure") + + class CacheRecorder: + def __init__(self, responses): + self.responses = responses + + def _sync_add_streaming_response_to_cache(self, response): + self.responses.append(response) + + class Observe(CustomLogger): + def __init__(self, responses, finished): + super().__init__() + self.responses = responses + self.finished = finished + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.responses.append(response_obj) + self.finished.set() + + try: + litellm.disable_streaming_logging = True + for fallback in (False, True): + cached, logged = [], [] + finished = threading.Event() + + logger = logger_for(stream=True, sync_callbacks=[Observe(logged, finished)]) + logger._llm_caching_handler = CacheRecorder(cached) + source = ControlledStream() + source.originals[1].usage = source.originals[2].usage + wrapper = CustomStreamWrapper( + completion_stream=iter(source.originals[:2]), + model="test", + logging_obj=logger, + custom_llm_provider="bedrock", + ) + pull = owners.prepare(wrapper.__next__, ()) + close = owners.prepare(wrapper.aclose, (), awaited=True) + try: + first = pull.invoke() + assert wrapper.chunks[0] is first + first.choices[0].delta.content = "consumer mutation" + shared = {"values": []} + last = pull.invoke() + assert last.choices[0].finish_reason == "stop" + for chunk in wrapper.chunks: + chunk._hidden_params["retained_shared"] = shared + if fallback: + chunk._hidden_params["uncopyable"] = Uncopyable() + retained_hidden = last._hidden_params + with TestCase().assertRaises(StopIteration): + pull.invoke() + assert await asyncio.to_thread(finished.wait, 5) + finally: + pull.close() + await close.invoke() + close.close() + assert len(cached) == len(logged) == 1 + cache_response, log_response = cached[0], logged[0] + assert cache_response is not log_response + assert cache_response.choices[0].message.content == "consumer mutation" + assert log_response.choices[0].message.content == "consumer mutation" + assert retained_hidden["usage"].total_tokens == 8 + assert (cache_response.choices is log_response.choices) is fallback + assert (cache_response.usage is log_response.usage) is fallback + assert (cache_response.usage is retained_hidden["usage"]) is fallback + assert (cache_response._hidden_params is log_response._hidden_params) is fallback + assert (cache_response._hidden_params["retained_shared"] is shared) is fallback + cache_response.choices[0].message.content = "cache mutation" + cache_response._hidden_params["retained_shared"]["values"].append("cache mutation") + assert log_response.choices[0].message.content == ("cache mutation" if fallback else "consumer mutation") + assert shared["values"] == (["cache mutation"] if fallback else []) + assert log_response._hidden_params["retained_shared"]["values"] == (["cache mutation"] if fallback else []) + assert copy_attempts == [True] + finally: + litellm.disable_streaming_logging = original_disable async def real_stream_close(owners): @@ -168,6 +453,7 @@ async def real_stream_close(owners): ) pull = owners.prepare(wrapper.__anext__, (), awaited=True) chunk = await pull.invoke() + assert wrapper.chunks[0] is chunk pull.close() retained = owners.prepare(lambda value: value, (chunk,)) close = owners.prepare(wrapper.aclose, (), awaited=True) @@ -175,10 +461,11 @@ async def real_stream_close(owners): await close.invoke() close.close() assert source.closed == 1 + assert wrapper.completion_stream is None + assert not getattr(logger, "_deferred_stream_complete_args", None) assert retained.invoke() is chunk assert chunk.choices[0].delta.content == "hello" retained.close() - await drain_component_tasks() async def real_stream_cancellation(owners): @@ -186,23 +473,51 @@ async def real_stream_cancellation(owners): class SuspendedStream(ControlledStream): async def __anext__(self): + if self.chunks is not None: + chunk = next(self.chunks) + self.chunks = None + return chunk entered.set() await asyncio.Event().wait() source = SuspendedStream() + source.originals[0].usage = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) logger = logger_for(stream=True) wrapper = CustomStreamWrapper( completion_stream=source, model="test", logging_obj=logger, custom_llm_provider="bedrock" ) pull = owners.prepare(wrapper.__anext__, (), awaited=True) + chunk = await pull.invoke() + retained = owners.prepare(lambda value: value, (chunk,)) + assert wrapper.chunks[0] is not chunk + assert wrapper.chunks[0].usage is source.originals[0].usage task = asyncio.create_task(pull.invoke()) pull.close() await entered.wait() task.cancel() with TestCase().assertRaises(asyncio.CancelledError): await task + assert logger.model_call_details.get("combined_usage_object") is None + recover = owners.prepare(wrapper._record_partial_usage_for_failure, ()) + try: + assert recover.invoke() is None + finally: + recover.close() + usage = logger.model_call_details["combined_usage_object"] + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (3, 2, 5) + assert usage is not source.originals[0].usage + source.originals[0].usage.total_tokens = 99 + assert usage.total_tokens == 5 + retained_usage = owners.prepare(lambda value: value, (usage,)) close = owners.prepare(wrapper.aclose, (), awaited=True) await close.invoke() + await close.invoke() close.close() assert source.closed == 1 - await drain_component_tasks() + assert wrapper.completion_stream is None and len(wrapper.chunks) == 1 + assert not getattr(logger, "_deferred_stream_complete_args", None) + assert retained.invoke() is chunk and chunk.choices[0].delta.content == "hello" + retained.close() + del wrapper, logger, usage + assert retained_usage.invoke().total_tokens == 5 + retained_usage.close() diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py new file mode 100644 index 00000000000..27aac2fbc59 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py @@ -0,0 +1,506 @@ +import asyncio +import copy +import gzip +import json +import threading +from collections import OrderedDict +from datetime import datetime +from functools import wraps +from unittest import TestCase + +import httpx + +import litellm +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from litellm.integrations.literal_ai import LiteralAILogger +from litellm.integrations.rubrik import RubrikLogger +from litellm.litellm_core_utils.litellm_logging import Logging, create_dummy_standard_logging_payload +from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import CrowdStrikeAIDRHandler +from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import MicrosoftPurviewDLPGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ModelResponse + + +async def integration_invoke(owners, callback, *args, **kwargs): + owner = owners.prepare(callback, args, kwargs, awaited=True) + pending = owner.invoke() + owner.close() + return await pending + + +def integration_response(url, body, status=200, headers=None): + return httpx.Response(status, json=body, headers=headers, request=httpx.Request("POST", url)) + + +def integration_callback_scope(scenario): + @wraps(scenario) + async def run(owners): + callbacks = tuple(litellm.callbacks) + try: + return await scenario(owners) + finally: + litellm.callbacks[:] = callbacks + + return run + + +@integration_callback_scope +async def real_logging_queue_chain(owners): + entered, release = asyncio.Event(), asyncio.Event() + uploads = [] + + class VertexTransport: + async def _ensure_access_token_async(self, **kwargs): + entered.set() + await release.wait() + return "fixture-token", "fixture-project" + + def _get_token_and_url(self, **kwargs): + return kwargs["auth_header"], None + + class Transport: + async def post(self, url, **kwargs): + wire = {**kwargs, "json": json.loads(json.dumps(kwargs["json"]))} if "json" in kwargs else kwargs + uploads.append((url, wire)) + return integration_response(url, {}, 202 if "datadog" in url else 200) + + datadog = DataDogLogger.__new__(DataDogLogger) + CustomBatchLogger.__init__(datadog, batch_size=100, flush_lock=asyncio.Lock()) + datadog.intake_url, datadog.DD_API_KEY, datadog.is_mock_mode = "https://datadog.invalid/logs", "test", False + datadog.async_client = Transport() + gcs = GCSBucketLogger.__new__(GCSBucketLogger) + CustomBatchLogger.__init__(gcs, batch_size=100) + gcs.log_queue = asyncio.Queue() + gcs.BUCKET_NAME, gcs.path_service_account_json = "fixture-bucket", None + gcs.vertex_instances = {"IAM_AUTH": VertexTransport()} + gcs.use_batched_logging = True + gcs.async_httpx_client = Transport() + literal = LiteralAILogger.__new__(LiteralAILogger) + CustomBatchLogger.__init__(literal, batch_size=100, flush_lock=asyncio.Lock()) + literal.literalai_api_url, literal.headers = "https://literal.invalid", {} + literal.async_httpx_client = Transport() + + payload = create_dummy_standard_logging_payload() + payload.update(status="failure", error_str="x" * 10001) + messages, settings, metadata = payload["messages"], payload["model_parameters"], payload["metadata"] + completion = payload["response"]["choices"][0]["message"] + tools = [{"type": "function", "function": {"name": "lookup"}}] + settings["tools"] = tools + now = datetime.now() + logging = Logging( + model="fixture-model", + messages=messages, + stream=False, + call_type="acompletion", + start_time=now, + litellm_call_id="fixture-queue", + function_id="fixture", + dynamic_async_failure_callbacks=[datadog, gcs, literal], + ) + error = RuntimeError("fixture failure") + kwargs = logging.model_call_details + kwargs.update(standard_logging_object=payload, model="fixture-model", exception=error, end_time=now) + await integration_invoke(owners, logging.async_failure_handler, error, "fixture traceback", now, now) + assert kwargs["standard_logging_object"] is payload + assert payload["messages"] is messages and payload["model_parameters"] is settings + assert payload["error_str"].endswith("truncated by litellm, this logger does not support large content") + assert "tools" not in settings + dd_snapshot = json.loads(datadog.log_queue[0]["message"]) + assert dd_snapshot["model_parameters"]["tools"] == tools + queued = gcs.log_queue.get_nowait() + assert queued["payload"] is payload and queued["kwargs"] is kwargs and queued["response_obj"] is None + gcs.log_queue.put_nowait(queued) + generation = literal.log_queue[0]["generation"] + assert generation["settings"] is settings and generation["tools"] is tools + assert generation["messages"] is messages and generation["messageCompletion"] is completion + assert literal.log_queue[0]["metadata"] is metadata + + flush = asyncio.create_task(integration_invoke(owners, gcs.flush_queue)) + try: + await entered.wait() + assert not uploads and not flush.done() and gcs.log_queue.empty() + messages[0]["content"] = "mutated before serialization" + settings["temperature"] = 0.25 + completion["content"] = "late completion" + payload["messages"] = [{"role": "user", "content": "replacement field"}] + kwargs["standard_logging_object"] = {"replacement": True} + release.set() + await flush + finally: + release.set() + if not flush.done(): + flush.cancel() + await asyncio.gather(flush, return_exceptions=True) + assert len(uploads) == 1 + gcs_snapshot = json.loads(uploads[0][1]["data"]) + assert gcs_snapshot["messages"] == payload["messages"] + assert gcs_snapshot["model_parameters"] == settings + assert "replacement" not in gcs_snapshot + await integration_invoke(owners, datadog.flush_queue) + await integration_invoke(owners, literal.flush_queue) + assert len(uploads) == 3 and not datadog.log_queue and not literal.log_queue + sent_dd = json.loads(gzip.decompress(uploads[1][1]["data"])) + assert json.loads(sent_dd[0]["message"]) == dd_snapshot + literal_wire = uploads[2][1]["json"] + sent_generation = literal_wire["variables"]["generation_0"] + assert sent_generation["messages"] == messages + assert sent_generation["messages"] != gcs_snapshot["messages"] + assert sent_generation["settings"]["temperature"] == 0.25 + assert sent_generation["messageCompletion"]["content"] == "late completion" + messages[0]["content"] = "after serialization" + assert sent_generation["messages"][0]["content"] == "mutated before serialization" + assert dd_snapshot["messages"][0]["content"] == "Hello, world!" + + +@integration_callback_scope +async def real_crowdstrike_translator_identity(owners): + entered, release = asyncio.Event(), asyncio.Event() + calls = [] + + class Transport: + async def post(self, url, json, **kwargs): + calls.append(json) + entered.set() + await release.wait() + return integration_response( + url, + { + "result": { + "blocked": False, + "transformed": True, + "guard_output": {"messages": [{"role": "user", "content": "redacted"}]}, + } + }, + ) + + guardrail = CrowdStrikeAIDRHandler.__new__(CrowdStrikeAIDRHandler) + CustomGuardrail.__init__(guardrail, guardrail_name="fixture-crowdstrike", event_hook=GuardrailEventHooks.pre_call) + guardrail.api_base, guardrail.api_key, guardrail.fail_on_error = "https://crowdstrike.invalid", "test", True + guardrail.skip_system_message_in_guardrail = True + guardrail.async_handler = Transport() + system = {"role": "system", "content": "internal policy"} + user = {"role": "user", "content": "private text", "extra": {"retained": True}} + messages = [system, user] + data = {"model": "fixture-model", "messages": messages} + task = asyncio.create_task( + integration_invoke( + owners, + OpenAIChatCompletionsHandler().process_input_messages, + data, + guardrail, + ) + ) + try: + await entered.wait() + assert data["messages"] is messages and not task.done() + assert calls[0]["guard_input"]["messages"] == [{"role": "user", "content": "private text"}] + user["extra"]["during_http"] = True + release.set() + assert await task is data + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert data["messages"] is not messages + assert data["messages"][0] is system + assert data["messages"][1] is not user + assert data["messages"][1]["extra"] is user["extra"] + assert data["messages"][1]["content"] == "redacted" + assert user["content"] == "private text" and messages[1] is user + detached = copy.deepcopy(messages[1:]) + inputs = {"texts": ["private text"], "structured_messages": detached} + control = await integration_invoke(owners, guardrail.apply_guardrail, inputs, {"messages": messages}, "request") + assert len(calls) == 2 + assert control["structured_messages"] is detached and detached[0] is not user + assert control["texts"] == ["redacted"] + assert detached[0]["content"] == user["content"] == "private text" + + +@integration_callback_scope +async def real_rubrik_block_lifecycle(owners): + for input_type, populated in (("request", False), ("response", True)): + entered, release = asyncio.Event(), asyncio.Event() + moderation, uploads = [], [] + + class Transport: + def __init__(self, moderation, uploads, entered, release): + self.moderation, self.uploads = moderation, uploads + self.entered, self.release = entered, release + + async def post(self, url, json, **kwargs): + if url.endswith("/batch"): + self.uploads.append(json) + return integration_response(url, {}) + self.moderation.append(json) + self.entered.set() + await self.release.wait() + return integration_response(url, {"choices": [{"message": {"content": "blocked by policy"}}]}) + + rubrik = RubrikLogger.__new__(RubrikLogger) + CustomGuardrail.__init__( + rubrik, + guardrail_name="fixture-rubrik", + event_hook=GuardrailEventHooks.post_call, + flush_lock=asyncio.Lock(), + batch_size=100, + ) + rubrik._periodic_flush_task = None + rubrik.sampling_rate, rubrik._headers = 1.0, {} + rubrik._dropped_since_warning, rubrik._last_drop_warning_time = 0, 0.0 + rubrik.prompt_moderation_endpoint = "https://rubrik.invalid/before" + rubrik.response_moderation_endpoint = "https://rubrik.invalid/after" + rubrik.logging_endpoint = "https://rubrik.invalid/batch" + rubrik.moderation_client = rubrik.async_httpx_client = Transport(moderation, uploads, entered, release) + other = RubrikLogger.__new__(RubrikLogger) + CustomGuardrail.__init__(other, guardrail_name="fixture-other-rubrik", event_hook=GuardrailEventHooks.post_call) + logging = Logging( + model="fixture-model", + messages=[{"role": "user", "content": "original prompt"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="fixture-correlation", + function_id="fixture", + ) + details = logging.model_call_details + details.update(messages=logging.messages, model="fixture-model", litellm_call_id="fixture-correlation") + details["system"] = "system scaffold" + if populated: + details["standard_logging_object"] = create_dummy_standard_logging_payload() + messages = details["messages"] + request = {"model": "fixture-model", "litellm_call_id": "fixture-correlation", "messages": messages} + inputs = {"texts": ["original response"], "structured_messages": messages} + success = owners.prepare(rubrik.async_log_success_event, (details, None, None, None), awaited=True) + task = asyncio.create_task( + integration_invoke(owners, rubrik.apply_guardrail, inputs, request, input_type, logging) + ) + try: + await entered.wait() + assert not task.done() and "_rubrik_logging_obj" not in request + assert "_rubrik_blocked" not in details + if input_type == "request": + assert moderation[0]["correlation_key"] == "fixture-correlation" + assert moderation[0]["messages"][0]["content"] == "original prompt" + else: + assert moderation[0]["request"]["messages"] is messages + assert moderation[0]["response"]["id"] == "fixture-correlation" + release.set() + with TestCase().assertRaises(ModifyResponseException) as caught: + await task + error = caught.exception + assert error.request_data is request and error.message == "blocked by policy" + assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True + await integration_invoke( + owners, other.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user") + ) + assert request["_rubrik_logging_obj"] is logging and details["_rubrik_blocked"] is True + assert not other.log_queue and not rubrik.log_queue and not uploads + await integration_invoke( + owners, rubrik.async_post_call_failure_hook, request, error, UserAPIKeyAuth(user_id="fixture-user") + ) + assert "_rubrik_logging_obj" not in request and details["_rubrik_blocked"] is True + assert len(rubrik.log_queue) == 1 + queued = rubrik.log_queue[0] + assert queued["id"] == "fixture-correlation" + assert queued["response"] == "ModifyResponseException: blocked by policy" + assert queued["messages"][0] == {"role": "system", "content": "system scaffold"} + assert messages[0] == {"role": "user", "content": "original prompt"} + if populated: + base = details["standard_logging_object"] + assert queued["metadata"] is not base["metadata"] + assert queued["messages"][1] is not base["messages"][0] + assert isinstance(base["response"], dict) + else: + assert queued["messages"][1] is messages[0] + assert queued["status"] == "failure" + assert queued["metadata"]["user_api_key_user_id"] == "fixture-user" + pending = success.invoke() + success.close() + assert await pending is None + assert len(rubrik.log_queue) == 1 and rubrik.log_queue[0] is queued + await integration_invoke(owners, rubrik.flush_queue) + assert not rubrik.log_queue and len(uploads) == 1 and uploads[0][0] is queued + finally: + success.close() + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await rubrik.aclose() + if rubrik._periodic_flush_task is not None: + await asyncio.gather(rubrik._periodic_flush_task, return_exceptions=True) + + +@integration_callback_scope +async def real_parallel_guardrail_snapshots(owners): + original_mode = litellm.safe_memory_mode + try: + for safe_memory_mode in (False, True): + litellm.safe_memory_mode = safe_memory_mode + await integration_parallel_snapshot_case(owners) + finally: + litellm.safe_memory_mode = original_mode + + +async def integration_parallel_snapshot_case(owners): + from litellm.caching.dual_cache import DualCache + from litellm.litellm_core_utils.core_helpers import independent_snapshot + from litellm.proxy.utils import ProxyLogging + + class Uncopyable: + def __init__(self): + self.attempts = 0 + self.observed = [] + + def __deepcopy__(self, memo): + self.attempts += 1 + raise TypeError("fixture cannot be copied") + + arrived, release, mutated = asyncio.Event(), asyncio.Event(), asyncio.Event() + observations = {} + sentinel = Uncopyable() + live = {"messages": [{"role": "user", "content": "original"}], "uncopyable": sentinel} + raw = independent_snapshot(live) + assert raw is not live and raw["messages"][0] is not live["messages"][0] + assert raw["uncopyable"] is sentinel and sentinel.attempts == 1 + live["messages"][0]["content"] = "masked" + + class Inspect(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + observations[self.guardrail_name] = data + if len(observations) == 4: + arrived.set() + await release.wait() + if self.guardrail_name == "fixture-live-writer": + data["messages"][0]["content"] = "shared mutation" + mutated.set() + await mutated.wait() + if self.scan_raw_request: + assert data["messages"][0]["content"] == "original" + data["messages"][0]["content"] = self.guardrail_name + assert data["uncopyable"] is sentinel + data["uncopyable"].observed.append(self.guardrail_name) + else: + assert data is live and data["messages"][0]["content"] == "shared mutation" + return {"discarded": self.guardrail_name} + + guardrails = tuple( + Inspect( + guardrail_name=name, + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + scan_raw_request="raw" in name, + ) + for name in ("fixture-live-writer", "fixture-live-reader", "fixture-raw-a", "fixture-raw-b") + ) + proxy = ProxyLogging.__new__(ProxyLogging) + proxy.call_details = {"user_api_key_cache": DualCache()} + task = asyncio.create_task( + integration_invoke( + owners, proxy._run_parallel_pre_call_guardrails, guardrails, live, raw, UserAPIKeyAuth(), "acompletion" + ) + ) + try: + await arrived.wait() + assert not task.done() + assert observations["fixture-live-writer"] is observations["fixture-live-reader"] is live + first, second = observations["fixture-raw-a"], observations["fixture-raw-b"] + assert first is not second and first is not raw and second is not raw + assert first["messages"][0] is not second["messages"][0] + assert first["messages"][0] is not raw["messages"][0] + assert first["uncopyable"] is second["uncopyable"] is raw["uncopyable"] is sentinel + assert sentinel.attempts == 3 and not sentinel.observed + release.set() + assert await task is None + assert raw["messages"] == [{"role": "user", "content": "original"}] + assert live["messages"][0]["content"] == "shared mutation" + assert set(sentinel.observed) == {"fixture-raw-a", "fixture-raw-b"} and len(sentinel.observed) == 2 + assert "discarded" not in live + assert first["messages"][0]["content"] == "fixture-raw-a" + assert second["messages"][0]["content"] == "fixture-raw-b" + assert all(guardrail._pre_call_hook_already_ran(live) for guardrail in guardrails if guardrail.scan_raw_request) + finally: + release.set() + mutated.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@integration_callback_scope +async def real_purview_sync_background(owners): + entered, release = threading.Event(), threading.Event() + calls, workers = [], [] + main_thread = threading.get_ident() + + class Transport: + async def post(self, url, **kwargs): + workers.append(threading.current_thread()) + calls.append((url, kwargs)) + if url.endswith("/token"): + entered.set() + assert release.wait(5), "background audit was not released" + return integration_response(url, {"access_token": "fixture-token", "expires_in": 3600}) + assert kwargs["headers"]["Authorization"] == "Bearer fixture-token" + if url.endswith("/compute"): + return integration_response(url, {}, headers={"etag": "fixture-etag"}) + assert url.endswith("/processContent") + assert kwargs["headers"]["If-None-Match"] == "fixture-etag" + return integration_response( + url, {"policyActions": [{"action": "restrictAccess", "restrictionAction": "block"}]} + ) + + purview = MicrosoftPurviewDLPGuardrail.__new__(MicrosoftPurviewDLPGuardrail) + CustomGuardrail.__init__(purview, guardrail_name="fixture-purview", event_hook=GuardrailEventHooks.logging_only) + purview.async_handler = Transport() + purview.tenant_id, purview.client_id, purview.client_secret = "fixture-tenant", "fixture-client", "test" + purview.purview_app_name, purview.user_id_field, purview.guardrail_provider = ( + "fixture", + "user_id", + "microsoft_purview", + ) + purview._token_cache, purview._scope_cache = None, OrderedDict() + purview._scope_cache_maxsize, purview._cache_lock = 1000, threading.Lock() + metadata = {"user_api_key_user_id": "fixture-user"} + kwargs = { + "messages": [{"role": "user", "content": "prompt at dispatch"}], + "litellm_params": {"metadata": metadata}, + "litellm_call_id": "before-http", + } + result = ModelResponse(model="fixture-model", choices=[{"message": {"role": "assistant", "content": "before"}}]) + owner = owners.prepare(purview.logging_hook, (kwargs, result, "completion"), awaited=False) + try: + returned = await asyncio.to_thread(owner.invoke) + owner.close() + assert returned[0] is kwargs and returned[1] is result + assert await asyncio.to_thread(entered.wait, 5) + assert len(calls) == 1 and workers[0].ident != main_thread and workers[0].daemon + assert workers[0].is_alive() + kwargs["messages"][0]["content"] = "too late for prompt extraction" + kwargs["litellm_call_id"] = "after-http" + result.choices[0].message.content = "response mutated while audit waits" + release.set() + await asyncio.to_thread(workers[0].join, 5) + assert not workers[0].is_alive() and all(worker is workers[0] for worker in workers) + assert len(calls) == 4 + entries = [call[1]["json"]["contentToProcess"] for call in calls[2:]] + assert [entry["activityMetadata"]["activity"] for entry in entries] == ["uploadText", "downloadText"] + assert entries[0]["contentEntries"][0]["content"]["data"] == "prompt at dispatch" + assert entries[1]["contentEntries"][0]["content"]["data"] == "response mutated while audit waits" + assert all(entry["contentEntries"][0]["correlationId"] == "after-http" for entry in entries) + assert kwargs["litellm_params"]["metadata"] is metadata + info = kwargs["metadata"]["standard_logging_guardrail_information"] + assert len(info) == 2 and all(item["guardrail_status"] == "guardrail_intervened" for item in info) + assert all(item["end_time"] >= item["start_time"] and item["duration"] >= 0 for item in info) + finally: + owner.close() + release.set() + if workers: + await asyncio.to_thread(workers[0].join, 5) 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 6a7da52b056..d9601f429a2 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py @@ -498,12 +498,107 @@ async def repeated_ownership(owners): assert all(ref() is None for ref in refs) +async def retained_field_replacement(owners): + original = {"messages": [{"content": "original"}]} + replacement = {"messages": [{"content": "replacement"}]} + event = {"payload": original, "alias": original} + saved = [] + + def retain(value): + saved.append(value["payload"]) + + def replace(value): + value["payload"] = replacement + value["alias"]["messages"][0]["content"] = "mutated original" + + for callback in (retain, replace): + owner = owners.prepare(callback, (event,)) + try: + assert owner.invoke() is None + finally: + owner.close() + assert saved[0] is original is event["alias"] + assert event["payload"] is replacement + assert saved[0]["messages"][0]["content"] == "mutated original" + replacement["messages"][0]["content"] = "mutated replacement" + assert event["payload"]["messages"][0]["content"] == "mutated replacement" + assert original["messages"][0]["content"] == "mutated original" + + +async def queued_graph_ownership(owners): + queue = asyncio.Queue() + sentinel = Value() + reference = weakref.ref(sentinel) + payload = {"sentinel": sentinel, "nested": {"status": "queued"}} + snapshot = json.dumps(payload["nested"]) + enqueue = owners.prepare(queue.put_nowait, (payload,)) + try: + enqueue.invoke() + finally: + enqueue.close() + del sentinel, payload + gc.collect() + assert owners.live == 0 and reference() is not None + queued = queue.get_nowait() + queued["nested"]["status"] = "changed before flush" + assert json.loads(json.dumps(queued["nested"])) == {"status": "changed before flush"} + assert json.loads(snapshot) == {"status": "queued"} + assert queued["sentinel"] is reference() + queue.task_done() + del queued + gc.collect() + assert reference() is None + + +async def detached_work_after_error(owners): + for raises in (False, True): + entered, release = asyncio.Event(), asyncio.Event() + tasks, observed = [], [] + value = Value() + value.status = "before return" + reference = weakref.ref(value) + + async def consume(argument, entered=entered, release=release, observed=observed): + entered.set() + await release.wait() + observed.append(argument.status) + + def callback(argument, tasks=tasks, consume=consume, raises=raises): + tasks.append(asyncio.create_task(consume(argument))) + if raises: + raise ValueError("after task creation") + + owner = owners.prepare(callback, (value,)) + try: + if raises: + with TestCase().assertRaisesRegex(ValueError, "after task creation"): + owner.invoke() + else: + assert owner.invoke() is None + finally: + owner.close() + del value + try: + await entered.wait() + assert owners.live == 0 and reference() is not None + reference().status = "after invocation" + release.set() + await tasks[0] + assert observed == ["after invocation"] + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + tasks.clear() + await checkpoint() + gc.collect() + assert reference() is None + + def run_scenario(name, retained, factory): owners = factory if retained else ReferenceFactory() async def run(): - async with asyncio.timeout(15): - await globals()[name](owners) + await asyncio.wait_for(globals()[name](owners), timeout=15) assert owners.live == 0 pending = asyncio.all_tasks() - {asyncio.current_task()} assert not pending, f"undrained tasks: {pending}" diff --git a/litellm-rust/crates/python-interop/tests/interop.rs b/litellm-rust/crates/python-interop/tests/interop.rs index 9c456dcb938..a40a73b2b6d 100644 --- a/litellm-rust/crates/python-interop/tests/interop.rs +++ b/litellm-rust/crates/python-interop/tests/interop.rs @@ -1,26 +1,12 @@ -use pyo3::Python; -use rstest::{fixture, rstest}; +use rstest::rstest; use serde_json::{Value, json}; use litellm_python_interop::{from_py, release_count, release_gil, to_py}; -struct InitializedPython; +#[path = "support/mod.rs"] +mod support; -impl InitializedPython { - fn attach(&self, f: F) -> R - where - F: for<'py> FnOnce(Python<'py>) -> R, - { - Python::attach(f) - } -} - -#[fixture] -#[once] -fn initialized_python() -> InitializedPython { - Python::initialize(); - InitializedPython -} +use support::python::{InitializedPython, initialized_python}; #[rstest] fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &InitializedPython) { diff --git a/litellm-rust/crates/python-interop/tests/prepared_call.rs b/litellm-rust/crates/python-interop/tests/prepared_call.rs index 100cb6383e5..a8549720c9d 100644 --- a/litellm-rust/crates/python-interop/tests/prepared_call.rs +++ b/litellm-rust/crates/python-interop/tests/prepared_call.rs @@ -1,19 +1,12 @@ -use std::ffi::CStr; - use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyTuple}; -use rstest::{fixture, rstest}; +use rstest::rstest; -fn scope<'py>(py: Python<'py>, source: &CStr) -> PyResult> { - let globals = PyDict::new(py); - py.run(source, Some(&globals), None)?; - Ok(globals) -} +#[path = "support/mod.rs"] +mod support; -fn item<'py>(globals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - globals.get_item(name).unwrap().unwrap() -} +use support::python::{InitializedPython, initialized_python, item, scope}; #[rstest] fn retains_aliases_mutations_and_original_result( @@ -366,15 +359,6 @@ assert first.body['document']['value'] == 'after invocation' }) } -struct InitializedPython; - -#[fixture] -#[once] -fn initialized_python() -> InitializedPython { - Python::initialize(); - InitializedPython -} - fn invoke_direct(call: &PreparedCall, py: Python<'_>) -> PyResult> { match call.invoke(py)? { InvocationOutcome::Returned(value) => Ok(value), diff --git a/litellm-rust/crates/python-interop/tests/support/mod.rs b/litellm-rust/crates/python-interop/tests/support/mod.rs new file mode 100644 index 00000000000..c0a088c7544 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/support/mod.rs @@ -0,0 +1 @@ +pub mod python; diff --git a/litellm-rust/crates/python-interop/tests/support/python.rs b/litellm-rust/crates/python-interop/tests/support/python.rs new file mode 100644 index 00000000000..5ef9c7e7b37 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/support/python.rs @@ -0,0 +1,47 @@ +#![allow(dead_code)] + +use std::ffi::CStr; + +use pyo3::prelude::*; +use pyo3::types::PyDict; +use rstest::fixture; + +pub struct InitializedPython; + +impl InitializedPython { + pub fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } +} + +#[fixture] +#[once] +pub fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython +} + +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(()) +} + +pub fn scope<'py>(py: Python<'py>, source: &CStr) -> PyResult> { + let globals = PyDict::new(py); + py.run(source, Some(&globals), None)?; + Ok(globals) +} + +pub fn item<'py>(globals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + globals.get_item(name).unwrap().unwrap() +} 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 ee6c5069910..a6e330afa26 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -36,14 +36,24 @@ class NativeRouteHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: content_length: Final = int(self.headers.get("content-length", "0")) - body: Final = json.loads(self.rfile.read(content_length)) + wire_body: Final = self.rfile.read(content_length) + body: Final = json.loads(wire_body) 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: - PUBLIC_OCR_REQUESTS.put(public_case) 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) @@ -266,6 +276,11 @@ def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: 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: @@ -273,26 +288,60 @@ def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: if event == "call" and frame.f_code is OCRRetainedBoundary.encode.__code__: observed["encode"] += 1 - class MutatingLogger(CustomLogger): - calls = 0 + 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]) -> None: + 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) - 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 headers["x-test-callback"] == "before-callback" - headers["x-test-callback"] = case - body["include_image_base64"] = True + 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()}} - callback: Final = MutatingLogger() + callbacks: Final = tuple(RetainedLogger(phase) for phase in ("retain", "mutate", "observe")) kwargs: Final = { "model": "mistral/mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "document": document, "api_base": api_base, "api_key": "sk-native", "extra_headers": { @@ -302,7 +351,7 @@ def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: "x-test-callback": "before-callback", }, "include_image_base64": False, - "callbacks": [callback], + "callbacks": list(callbacks), "rust": True, "timeout": 3.0, "num_retries": 0, @@ -316,8 +365,26 @@ def exercise_public_ocr(install_root: Path, api_base: str, case: str) -> int: assert isinstance(response, OCRResponse) assert_success("ocr", response.model_dump()) assert response.model == "mistral-ocr-latest" - assert callback.calls == 1, callback.calls + 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