From 9de4e84feb128f2248253d9f52ae13049f2c88a5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:07:57 -0700 Subject: [PATCH] refactor(rust): extract domain-neutral Python interop --- litellm-rust/AGENTS.md | 9 +- litellm-rust/CLAUDE.md | 9 +- litellm-rust/Cargo.lock | 109 +++++++++++++++++- litellm-rust/Cargo.toml | 3 + litellm-rust/README.md | 8 +- .../PROVIDER_CODING_STANDARDS.md | 2 +- litellm-rust/crates/ai-gateway/README.md | 7 +- .../core/tests/workspace_crate_allowlist.rs | 16 ++- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- litellm-rust/crates/python-bridge/CLAUDE.md | 5 +- litellm-rust/crates/python-bridge/Cargo.toml | 3 +- .../python-bridge/benches/serialization.rs | 11 +- litellm-rust/crates/python-bridge/src/gil.rs | 32 ----- litellm-rust/crates/python-bridge/src/lib.rs | 16 +-- .../python-bridge/tests/marshal_boundary.rs | 11 +- litellm-rust/crates/python-interop/AGENTS.md | 1 + litellm-rust/crates/python-interop/Cargo.toml | 15 +++ litellm-rust/crates/python-interop/src/gil.rs | 21 ++++ litellm-rust/crates/python-interop/src/lib.rs | 5 + .../src/marshal.rs | 0 .../crates/python-interop/tests/interop.rs | 44 +++++++ 21 files changed, 247 insertions(+), 82 deletions(-) delete mode 100644 litellm-rust/crates/python-bridge/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/AGENTS.md create mode 100644 litellm-rust/crates/python-interop/Cargo.toml create mode 100644 litellm-rust/crates/python-interop/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/src/lib.rs rename litellm-rust/crates/{python-bridge => python-interop}/src/marshal.rs (100%) create mode 100644 litellm-rust/crates/python-interop/tests/interop.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 36a5ad5a8f4..b8b6291283d 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,6 +1,6 @@ # AGENTS.md -litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers. +litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates @@ -8,9 +8,10 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes ( |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Where a route lives @@ -28,7 +29,7 @@ core/src/messages/ Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. -Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. +Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index fe6ceedbb86..3dcf1853efc 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -21,12 +21,13 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -## Crates (exactly three — see AGENTS.md) +## Crates (see AGENTS.md) `litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. `litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not -a route — add modules, not crates. +`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` +holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate +is a layer or shared foundation, not a route; add modules, not crates. ## Core Boundary @@ -175,7 +176,7 @@ cd litellm-rust cargo fmt --check # the ai-gateway binary + server code is behind the `server` feature cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings -cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings +cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4388e561026..dd41cf0e84b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -919,6 +919,12 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + [[package]] name = "futures-util" version = "0.3.33" @@ -972,6 +978,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + [[package]] name = "h2" version = "0.3.27" @@ -1432,14 +1444,24 @@ dependencies = [ "criterion", "litellm-ai-gateway", "litellm-core", + "litellm-python-interop", "pyo3", "pyo3-async-runtimes", - "pythonize", - "serde", "serde_json", "tokio", ] +[[package]] +name = "litellm-python-interop" +version = "0.1.0" +dependencies = [ + "pyo3", + "pythonize", + "rstest", + "serde", + "serde_json", +] + [[package]] name = "litemap" version = "0.8.2" @@ -1627,6 +1649,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1899,6 +1930,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + [[package]] name = "reqwest" version = "0.12.28" @@ -1956,6 +1993,35 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.119", + "unicode-ident", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2488,6 +2554,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -2903,6 +2999,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c17a0605fc7..c447d915abe 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -2,6 +2,7 @@ members = [ "crates/core", "crates/ai-gateway", + "crates/python-interop", "crates/python-bridge", ] resolver = "2" @@ -15,12 +16,14 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } +litellm-python-interop = { path = "crates/python-interop" } 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"] } +rstest = "0.26.1" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index bcccf93300b..a0d79c6f0a5 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -26,9 +26,10 @@ coverage and production evidence. |-------|------| | litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | | litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Layout @@ -38,7 +39,8 @@ crates/ src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client src/providers/anthropic/messages/transformation.rs ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. - python-bridge/ PyO3 bridge for Python LiteLLM. + python-interop/ Domain-neutral PyO3 conversion and GIL primitives. + python-bridge/ PyO3 API adapter for Python LiteLLM. ``` The folder shape follows the Python provider tree: diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index a1860d8a9c9..4a689cb9579 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -54,6 +54,6 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` cd litellm-rust cargo fmt --check cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings - cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings + cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 7a6c620ee84..5cbb47220be 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,15 +6,16 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` is exactly three crates (a crate is a **layer**, not a route): +`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index 656ba033b62..8a8a5ea263a 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,7 +1,8 @@ -//! Enforcement: the litellm-rust workspace has exactly three crates. +//! Enforcement: the litellm-rust workspace has exactly four crates. //! -//! `core` (pure translation), `ai-gateway` (routes + all network I/O), and -//! `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a +//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host), +//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the +//! PyO3 cdylib). Adding or removing a crate must be a //! deliberate act: this test fails until the allowlist here is updated, forcing //! whoever changes the crate set to justify the new crate per the rule that a //! crate is a layer needing independent compilation / its own deps / a separate @@ -16,10 +17,15 @@ use std::path::{Path, PathBuf}; /// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the /// workspace legitimately gains or loses a crate. -const EXPECTED_MEMBERS: &[&str] = &["crates/core", "crates/ai-gateway", "crates/python-bridge"]; +const EXPECTED_MEMBERS: &[&str] = &[ + "crates/core", + "crates/ai-gateway", + "crates/python-interop", + "crates/python-bridge", +]; /// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-bridge"]; +const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"]; const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index ad3cddfa5fd..42282ca4da4 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,3 @@ -litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`). +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. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index 3ce8b8c639a..d25ae5a8130 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -5,8 +5,9 @@ Rules for `litellm-rust/crates/python-bridge`. ## Responsibility `python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. -Keep this crate thin. It adapts Python objects to Rust payloads and returns -Python-compatible dictionaries. +Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, +maps domain errors to Python exceptions, and delegates generic conversion and +GIL handling to `litellm-python-interop`. ## Bridge Shape diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index b1fdfd7677a..498003de149 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -18,10 +18,9 @@ panic-test = [] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } +litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true -pythonize.workspace = true -serde.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs index 8a90cf667d0..0b9436d0cb7 100644 --- a/litellm-rust/crates/python-bridge/benches/serialization.rs +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -2,6 +2,7 @@ use std::hint::black_box; use std::time::Duration; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use litellm_python_interop::{from_py, to_py}; use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Value, json}; @@ -25,7 +26,7 @@ fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Va } fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { - pythonize::depythonize(value).expect("payload should depythonize") + from_py(value).expect("payload should depythonize") } fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { @@ -37,12 +38,10 @@ fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { } fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { - pythonize::pythonize(py, value) - .expect("response should pythonize") - .unbind() + to_py(py, value).expect("response should pythonize") } -fn serialization(c: &mut Criterion) { +fn bridge_serialization(c: &mut Criterion) { Python::initialize(); Python::attach(|py| { for &(label, payload_bytes) in PAYLOAD_SIZES { @@ -98,6 +97,6 @@ criterion_group! { .sample_size(20) .warm_up_time(Duration::from_secs(1)) .measurement_time(Duration::from_secs(4)); - targets = serialization + targets = bridge_serialization } criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs deleted file mode 100644 index e887c8ec1e3..00000000000 --- a/litellm-rust/crates/python-bridge/src/gil.rs +++ /dev/null @@ -1,32 +0,0 @@ -//! GIL accounting. -//! -//! A single chokepoint for releasing the GIL around blocking work. Every -//! blocking call in the bridge goes through [`release_gil`] instead of calling -//! `Python::detach` directly, so the release count stays accurate and we -//! have one place to extend later (timing histograms, per-call labels, etc.). - -use std::sync::atomic::{AtomicU64, Ordering}; - -use pyo3::prelude::*; - -/// Number of times the bridge has released the GIL since process start. -static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); - -/// Release the GIL around `f`, recording the release. -/// -/// `f` must not touch any Python state — that is what makes releasing the GIL -/// safe. Returning the value back to Python re-acquires the GIL at the call -/// site, after `f` has finished. -pub fn release_gil(py: Python<'_>, f: F) -> T -where - F: FnOnce() -> T + Send, - T: Send, -{ - GIL_RELEASES.fetch_add(1, Ordering::Relaxed); - py.detach(f) -} - -/// Total GIL releases performed by the bridge so far. -pub fn release_count() -> u64 { - GIL_RELEASES.load(Ordering::Relaxed) -} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 18e0b05cbb5..746e0770f9b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,16 +13,12 @@ use litellm_core::chat_completions::{ use litellm_core::error::CoreError; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; -mod gil; -mod marshal; - -use marshal::{from_py, to_py}; - pyo3::create_exception!( _native, RustBridgeDeclined, @@ -230,7 +226,7 @@ fn ocr( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { model: &model, document, @@ -318,7 +314,7 @@ fn transcription( }; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; let timeout = optional_timeout(timeout_seconds); - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( AudioTranscriptionRequest { model: &model, @@ -419,7 +415,7 @@ fn messages( let (body, extra_headers, timeout) = marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { model: &model, body, @@ -546,7 +542,7 @@ fn chat_completions( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( ChatCompletionsRequest { model: &model, @@ -610,7 +606,7 @@ fn achat_completions( #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); - stats.set_item("releases", gil::release_count())?; + stats.set_item("releases", release_count())?; Ok(stats.into_any().unbind()) } diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs index 6a6ede22e85..d397d20b9fd 100644 --- a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -1,7 +1,7 @@ use std::fs; use std::path::{Path, PathBuf}; -const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ +const DISALLOWED_OUTSIDE_INTEROP: &[&str] = &[ "py.import(\"json\")", "pythonize::", "serde_json::to_string", @@ -33,18 +33,15 @@ fn rust_sources(directory: &Path) -> Vec { } #[test] -fn serialization_is_centralized_in_marshal_module() { +fn serialization_uses_the_interop_boundary() { 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 { + for disallowed in DISALLOWED_OUTSIDE_INTEROP { assert!( !source.contains(disallowed), - "{} bypasses the typed marshal module with `{disallowed}`", + "{} bypasses litellm-python-interop with `{disallowed}`", path.display() ); } diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md new file mode 100644 index 00000000000..d1d61e5dfa0 --- /dev/null +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -0,0 +1 @@ +litellm-python-interop is the domain-neutral PyO3 foundation. Keep generic Python/Serde conversion and interpreter primitives here. Do not add LiteLLM domain crates, route types, API registration, or cdylib build features. diff --git a/litellm-rust/crates/python-interop/Cargo.toml b/litellm-rust/crates/python-interop/Cargo.toml new file mode 100644 index 00000000000..9da6af6e2e2 --- /dev/null +++ b/litellm-rust/crates/python-interop/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-python-interop" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +pyo3.workspace = true +pythonize.workspace = true +serde.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/python-interop/src/gil.rs b/litellm-rust/crates/python-interop/src/gil.rs new file mode 100644 index 00000000000..04b966a6002 --- /dev/null +++ b/litellm-rust/crates/python-interop/src/gil.rs @@ -0,0 +1,21 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +use pyo3::prelude::*; + +static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// Runs work detached from the interpreter and records the release. +/// +/// `f` must not access Python state while the interpreter is detached. +pub fn release_gil(py: Python<'_>, f: F) -> T +where + F: FnOnce() -> T + Send, + T: Send, +{ + GIL_RELEASES.fetch_add(1, Ordering::Relaxed); + py.detach(f) +} + +pub fn release_count() -> u64 { + GIL_RELEASES.load(Ordering::Relaxed) +} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs new file mode 100644 index 00000000000..df2bd260fdb --- /dev/null +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -0,0 +1,5 @@ +mod gil; +mod marshal; + +pub use gil::{release_count, release_gil}; +pub use marshal::{from_py, to_py}; diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs similarity index 100% rename from litellm-rust/crates/python-bridge/src/marshal.rs rename to litellm-rust/crates/python-interop/src/marshal.rs diff --git a/litellm-rust/crates/python-interop/tests/interop.rs b/litellm-rust/crates/python-interop/tests/interop.rs new file mode 100644 index 00000000000..9c456dcb938 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/interop.rs @@ -0,0 +1,44 @@ +use pyo3::Python; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; + +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; + +struct InitializedPython; + +impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } +} + +#[fixture] +#[once] +fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython +} + +#[rstest] +fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { + let expected = json!({"model": "test", "items": [1, true, null]}); + let python_value = to_py(py, &expected).expect("value should convert to Python"); + let actual: Value = + from_py(python_value.bind(py)).expect("Python value should convert to serde"); + + assert_eq!(actual, expected); + }); +} + +#[rstest] +fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) { + let before = release_count(); + let result = python.attach(|py| release_gil(py, || 42)); + + assert_eq!(result, 42); + assert_eq!(release_count(), before + 1); +}