From 352789257d2732f93455fff914de9813e369ec50 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Sat, 29 Aug 2026 10:41:10 -0700 Subject: [PATCH] perf(rust): use pythonize for bridge serialization (#38764) * perf(rust): use pythonize for bridge serialization * perf(rust): benchmark serialization payload sizes * refactor(rust): remove mutable benchmark state --- litellm-rust/Cargo.lock | 340 ++++++++++++++++++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 15 +- .../python-bridge/benches/serialization.rs | 103 ++++++ litellm-rust/crates/python-bridge/src/lib.rs | 48 +-- .../crates/python-bridge/src/marshal.rs | 20 ++ .../python-bridge/tests/marshal_boundary.rs | 52 +++ 7 files changed, 547 insertions(+), 32 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/benches/serialization.rs create mode 100644 litellm-rust/crates/python-bridge/src/marshal.rs create mode 100644 litellm-rust/crates/python-bridge/tests/marshal_boundary.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce28f737334..4388e561026 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,36 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "arc-swap" version = "1.9.2" @@ -506,6 +536,12 @@ dependencies = [ "either", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.3.0" @@ -541,6 +577,58 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.58" @@ -596,6 +684,72 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -856,6 +1010,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -1179,6 +1344,15 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1255,10 +1429,13 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "criterion", "litellm-ai-gateway", "litellm-core", "pyo3", "pyo3-async-runtimes", + "pythonize", + "serde", "serde_json", "tokio", ] @@ -1340,6 +1517,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1352,6 +1535,16 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1376,6 +1569,34 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.14.0" @@ -1486,6 +1707,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "pythonize" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.11" @@ -1613,12 +1844,61 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + [[package]] name = "regex-lite" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1774,6 +2054,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -2099,6 +2388,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.12.0" @@ -2363,6 +2662,16 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2475,6 +2784,37 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-link" version = "0.2.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 6d63be05d00..481ea3f8f66 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.29.0" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } +pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 20a9ba789ce..0c4a753f762 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -9,10 +9,23 @@ repository.workspace = true name = "_native" crate-type = ["cdylib"] +[features] +default = ["extension-module"] +extension-module = ["pyo3/extension-module"] + [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } -pyo3 = { workspace = true, features = ["extension-module"] } +pyo3.workspace = true pyo3-async-runtimes.workspace = true +pythonize.workspace = true +serde.workspace = true serde_json.workspace = true tokio.workspace = true + +[dev-dependencies] +criterion = "0.8.2" + +[[bench]] +name = "serialization" +harness = false diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs new file mode 100644 index 00000000000..8a90cf667d0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -0,0 +1,103 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Value, json}; + +const PAYLOAD_SIZES: &[(&str, usize)] = &[ + ("1_KiB", 1024), + ("64_KiB", 64 * 1024), + ("1_MiB", 1024 * 1024), + ("4_MiB", 4 * 1024 * 1024), + ("16_MiB", 16 * 1024 * 1024), +]; + +fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value { + let json = py.import("json").expect("Python json module should import"); + let encoded: String = json + .call_method1("dumps", (value,)) + .expect("payload should serialize") + .extract() + .expect("json.dumps should return a string"); + serde_json::from_str(&encoded).expect("serialized JSON should parse") +} + +fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { + pythonize::depythonize(value).expect("payload should depythonize") +} + +fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { + let json = py.import("json").expect("Python json module should import"); + let encoded = serde_json::to_string(value).expect("response should serialize"); + json.call_method1("loads", (encoded,)) + .expect("serialized response should parse in Python") + .unbind() +} + +fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { + pythonize::pythonize(py, value) + .expect("response should pythonize") + .unbind() +} + +fn serialization(c: &mut Criterion) { + Python::initialize(); + Python::attach(|py| { + for &(label, payload_bytes) in PAYLOAD_SIZES { + let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes)); + let document = PyDict::new(py); + document + .set_item("type", "image_url") + .expect("document type should be set"); + document + .set_item("image_url", &data_uri) + .expect("document URL should be set"); + let response = json!({ + "pages": [{ + "index": 0, + "markdown": "OCR text", + "images": [{"image_base64": data_uri}], + }], + "model": "mistral-ocr-latest", + "document_annotation": null, + "usage_info": {"pages_processed": 1}, + "object": "ocr", + }); + + c.bench_with_input( + BenchmarkId::new("python_to_rust_json", label), + &document, + |b, document| { + b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any()))) + }, + ); + c.bench_with_input( + BenchmarkId::new("python_to_rust_pythonize", label), + &document, + |b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_json", label), + &response, + |b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_pythonize", label), + &response, + |b, response| b.iter(|| pythonize_to_py(py, black_box(response))), + ); + } + }); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(4)); + targets = serialization +} +criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c6f81cf6916..f9e75f45f75 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; mod gil; +mod marshal; + +use marshal::{from_py, to_py}; pyo3::create_exception!( _native, @@ -41,35 +44,18 @@ type MarshaledOcrInputs = ( Option, ); -fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { - let json = py.import("json")?; - let encoded: String = json.call_method1("dumps", (value,))?.extract()?; - serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) -} - -fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { - let json = py.import("json")?; - let encoded = - serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; - Ok(json.call_method1("loads", (encoded,))?.unbind()) -} - fn messages_response_to_py( py: Python<'_>, response: AnthropicMessagesResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn chat_completions_response_to_py( py: Python<'_>, response: ChatCompletionsResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn core_error_to_pyerr(err: CoreError) -> PyErr { @@ -116,7 +102,7 @@ fn optional_object_to_map( value: Option>, ) -> PyResult> { match value { - Some(value) => match py_to_json(py, value.bind(py))? { + Some(value) => match from_py(value.bind(py))? { Value::Object(map) => Ok(map), _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), }, @@ -139,7 +125,7 @@ fn marshal_headers( headers: Option>, ) -> PyResult> { let value = match headers { - Some(headers) => py_to_json(py, headers.bind(py))?, + Some(headers) => from_py(headers.bind(py))?, None => Value::Object(Map::new()), }; let Value::Object(headers) = value else { @@ -211,7 +197,7 @@ fn marshal_inputs( optional_params: Option>, timeout_seconds: Option, ) -> PyResult { - let document = py_to_json(py, document.bind(py))?; + let document = from_py(document.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -262,7 +248,7 @@ fn ocr( }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -307,7 +293,7 @@ fn aocr( .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -325,7 +311,7 @@ fn transcription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -351,7 +337,7 @@ fn transcription( )) }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -370,7 +356,7 @@ fn atranscription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -394,7 +380,7 @@ fn atranscription( }) .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -406,7 +392,7 @@ fn marshal_messages_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let body = py_to_json(py, body.bind(py))?; + let body: Value = from_py(body.bind(py))?; if !body.is_object() { return Err(PyValueError::new_err("body must be a dict")); } @@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let messages = py_to_json(py, messages.bind(py))?; + let messages: Value = from_py(messages.bind(py))?; if !messages.is_array() { return Err(PyValueError::new_err("messages must be a list")); } @@ -527,7 +513,7 @@ fn chat_completions_decline( optional_params: Option>, custom_llm_provider: Option, ) -> PyResult> { - let messages = py_to_json(py, messages.bind(py))?; + let messages = from_py(messages.bind(py))?; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; Ok(chat_completions_decline_reason( &model, diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..c3d0638427c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,20 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde::Serialize; +use serde::de::DeserializeOwned; + +pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) +} + +pub fn to_py(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(|error| PyValueError::new_err(error.to_string())) +} diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs new file mode 100644 index 00000000000..6a6ede22e85 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -0,0 +1,52 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ + "py.import(\"json\")", + "pythonize::", + "serde_json::to_string", + "serde_json::from_str", +]; + +fn source_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("src") +} + +fn rust_sources(directory: &Path) -> Vec { + fs::read_dir(directory) + .expect("bridge source directory should be readable") + .map(|entry| { + entry + .expect("bridge source entry should be readable") + .path() + }) + .flat_map(|path| { + if path.is_dir() { + rust_sources(&path) + } else if path.extension().is_some_and(|extension| extension == "rs") { + vec![path] + } else { + Vec::new() + } + }) + .collect() +} + +#[test] +fn serialization_is_centralized_in_marshal_module() { + let root = source_root(); + + for path in rust_sources(&root) { + if path == root.join("marshal.rs") { + continue; + } + let source = fs::read_to_string(&path).expect("bridge source should be readable"); + for disallowed in DISALLOWED_OUTSIDE_MARSHAL { + assert!( + !source.contains(disallowed), + "{} bypasses the typed marshal module with `{disallowed}`", + path.display() + ); + } + } +}