scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

wip

Browse source
This commit is contained in:
Yujong Lee 2026-09-06 18:43:05 -07:00
parent c6cd873266
commit 2ca2d106f7
34 changed files with 1843 additions and 585 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-rust.yml vendored
View file

@ -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

7
Makefile
View file

@ -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

8
litellm-rust/CLAUDE.md
View file

@ -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
```

71
litellm-rust/Cargo.lock generated
View file

@ -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"

2
litellm-rust/Cargo.toml
View file

@ -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"] }

53
litellm-rust/README.md
View file

@ -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

2
litellm-rust/crates/core/src/constants.rs
View file

@ -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";

16
litellm-rust/crates/core/src/ocr/transport.rs → litellm-rust/crates/core/src/http_utils/buffered_post.rs
View file

@ -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<Response, Error> {
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<Response, Error> {
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<Response, Error> {
.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<Response, Error> {
.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<Response, Error> {
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,

8
litellm-rust/crates/core/src/ocr/transport/tests.rs → litellm-rust/crates/core/src/http_utils/buffered_post/tests.rs
View file

@ -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")
);
}

2
litellm-rust/crates/core/src/http_utils.rs → litellm-rust/crates/core/src/http_utils/mod.rs
View file

@ -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;

1
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -1,3 +1,2 @@
pub mod transformation;
pub mod transport;
pub mod types;

3
litellm-rust/crates/python-bridge/AGENTS.md
View file

@ -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.

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -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

11
litellm-rust/crates/python-bridge/src/execution.rs
View file

@ -89,12 +89,19 @@ where
F: Future<Output = Result<T, Error>> + 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<T, F>(future: F, map_error: fn(Error) -> PyErr) -> PyResult<T>
where
T: Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
let result = catch_future_panic(future).await?;
map_core_result(result, map_error)
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),

39
litellm-rust/crates/python-bridge/src/routes/bindings.rs
View file

@ -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,
},
}
}
}

2
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -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)?;

132
litellm-rust/crates/python-bridge/src/routes/ocr_retained.rs
View file

@ -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<Py<PyAny>> {
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<Py<PyAny>> {
invoke(
boundary,
BoundaryMethod::Prepare.resolve(asynchronous),
PyTuple::empty(boundary.py()),
)
}
fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Request> {
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<Self::Output> {
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<Bound<'py, PyAny>> {
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<Py<PyAny>> {
invoke(
boundary,
BoundaryMethod::Finish.resolve(asynchronous),
PyTuple::new(boundary.py(), [wire])?,
)
}
use super::retained_http;
#[pyfunction]
fn ocr_retained(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
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<Py<PyAny>> {
static DRIVER: PyOnceLock<Py<PyAny>> = 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<()> {

172
litellm-rust/crates/python-bridge/src/routes/retained_http.rs Normal file
View file

@ -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<Py<PyAny>> {
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<Py<PyAny>> {
invoke(
boundary,
BoundaryMethod::Prepare.resolve(asynchronous),
PyTuple::empty(boundary.py()),
)
}
fn encode(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult<Request> {
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<Self::Output> {
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<Bound<'py, PyAny>> {
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<Py<PyAny>> {
invoke(
boundary,
BoundaryMethod::Finish.resolve(asynchronous),
PyTuple::new(boundary.py(), [wire])?,
)
}
pub(crate) fn run_sync(boundary: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
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<Py<PyAny>> {
static DRIVER: PyOnceLock<Py<PyAny>> = 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)?,
),
)
}

80
litellm-rust/crates/python-bridge/tests/fixtures/ocr_retained.py vendored
View file

@ -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())

269
litellm-rust/crates/python-bridge/tests/fixtures/retained_http_contract.py vendored Normal file
View file

@ -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")

268
litellm-rust/crates/python-bridge/tests/ocr_retained.rs
View file

@ -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),
)
})
}

47
litellm-rust/crates/python-bridge/tests/retained_http.rs Normal file
View file

@ -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(())
})
}

1
litellm-rust/crates/python-bridge/tests/support/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod native;

21
litellm-rust/crates/python-bridge/tests/support/native.rs Normal file
View file

@ -0,0 +1,21 @@
use pyo3::prelude::*;
use pyo3::types::PyDict;
pub fn native_globals(py: Python<'_>) -> PyResult<Bound<'_, PyDict>> {
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(())
}

1
litellm-rust/crates/python-interop/Cargo.toml
View file

@ -12,4 +12,5 @@ serde.workspace = true
[dev-dependencies]
rstest.workspace = true
serial_test.workspace = true
serde_json.workspace = true

70
litellm-rust/crates/python-interop/tests/callback_lifecycle.rs
View file

@ -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<PyDict> {
@ -65,6 +50,10 @@ fn scenario_scope(initialized_python: &InitializedPython) -> Py<PyDict> {
#[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<PyDict>,
#[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<PyDict>,
#[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<PyDict>,
#[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<PyDict>, #[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<PyDict>, #[case] awaited: bool)
}
#[rstest]
#[parallel(python_interpreter)]
fn awaited_raise_surfaces_when_driven(scenario_scope: Py<PyDict>) -> 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<PyDict>) -> PyResult<()
Some(globals),
None,
)?;
let started = || -> PyResult<usize> { Ok(globals.get_item("events")?.unwrap().len()?) };
let started = || -> PyResult<usize> { globals.get_item("events")?.unwrap().len() };
let call = PreparedCall::new(
InvocationMode::Await,
globals.get_item("callback")?.unwrap().unbind(),

339
litellm-rust/crates/python-interop/tests/fixtures/callback_components.py vendored
View file

@ -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()

506
litellm-rust/crates/python-interop/tests/fixtures/callback_integrations.py vendored Normal file
View file

@ -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)

99
litellm-rust/crates/python-interop/tests/fixtures/callback_lifecycle.py vendored
View file

@ -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}"

22
litellm-rust/crates/python-interop/tests/interop.rs
View file

@ -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<F, R>(&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) {

24
litellm-rust/crates/python-interop/tests/prepared_call.rs
View file

@ -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<Bound<'py, PyDict>> {
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<Py<PyAny>> {
match call.invoke(py)? {
InvocationOutcome::Returned(value) => Ok(value),

1
litellm-rust/crates/python-interop/tests/support/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod python;

47
litellm-rust/crates/python-interop/tests/support/python.rs Normal file
View file

@ -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<F, R>(&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<Bound<'py, PyDict>> {
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()
}

101
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

@ -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