From 2d77cc368c4067ff7910efa23f1d2323ae9f1db8 Mon Sep 17 00:00:00 2001 From: Jacinta Ferrant Date: Thu, 18 Jun 2026 20:53:44 -0700 Subject: [PATCH] feat(rust): add litellm_core PyO3 extension with regex pattern prefilter Signed-off-by: Jacinta Ferrant --- .gitignore | 1 + rust/Cargo.lock | 177 ++++++++++++++++++++++++++++++++++ rust/Cargo.toml | 12 +++ rust/pyproject.toml | 12 +++ rust/src/lib.rs | 16 +++ rust/src/pattern_prefilter.rs | 63 ++++++++++++ 6 files changed, 281 insertions(+) create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/pyproject.toml create mode 100644 rust/src/lib.rs create mode 100644 rust/src/pattern_prefilter.rs diff --git a/.gitignore b/.gitignore index fda3311fe02..e7f0579d50e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .python-version .venv +rust/target/ .venv_policy_test .env .claude diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 00000000000..28cf3ac4081 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,177 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "litellm-rs" +version = "0.1.0" +dependencies = [ + "pyo3", + "regex", +] + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "regex" +version = "1.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 00000000000..485ae13b032 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "litellm-rs" +version = "0.1.0" +edition = "2024" + +[lib] +name = "litellm_core" +crate-type = ["cdylib"] + +[dependencies] +pyo3 = { version = "0.29.0", features = ["extension-module"] } +regex = "1" diff --git a/rust/pyproject.toml b/rust/pyproject.toml new file mode 100644 index 00000000000..6567f2c068f --- /dev/null +++ b/rust/pyproject.toml @@ -0,0 +1,12 @@ +[build-system] +requires = ["maturin>=1.0,<2.0"] +build-backend = "maturin" + +[project] +name = "litellm-core" +version = "0.1.0" +requires-python = ">=3.9" + +[tool.maturin] +module-name = "litellm_core" +features = ["pyo3/extension-module"] diff --git a/rust/src/lib.rs b/rust/src/lib.rs new file mode 100644 index 00000000000..c3f9fc6fa20 --- /dev/null +++ b/rust/src/lib.rs @@ -0,0 +1,16 @@ +use pyo3::{ + Bound, PyResult, + types::{PyModule, PyModuleMethods}, + wrap_pyfunction, +}; + +mod pattern_prefilter; +use pattern_prefilter::build_pattern_prefilter; + +// Function name must match Cargo.toml's [lib] name and pyproject.toml's +// module-name: pyo3 derives the Python module name from it. +#[pyo3::pymodule] +fn litellm_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(build_pattern_prefilter, m)?)?; + Ok(()) +} diff --git a/rust/src/pattern_prefilter.rs b/rust/src/pattern_prefilter.rs new file mode 100644 index 00000000000..beb6e1354c5 --- /dev/null +++ b/rust/src/pattern_prefilter.rs @@ -0,0 +1,63 @@ +use pyo3::pyclass; +use pyo3::pymethods; +use regex::{RegexBuilder, RegexSet, RegexSetBuilder}; + +/// Conservative "could any of these patterns possibly match?" gate. +/// +/// Built from a fixed set of regex sources that each compiled successfully on +/// their own. `any_match` defaults to `true` (don't skip) whenever there is no +/// usable set, since the only safe direction to fail in is "check anyway." +#[pyclass] +pub struct PatternPrefilter { + regex_set: Option, +} + +#[pymethods] +impl PatternPrefilter { + fn any_match(&self, text: &str) -> bool { + self.regex_set.as_ref().is_none_or(|set| set.is_match(text)) + } +} + +/// Splits `patterns` into ones a `RegexSet` can evaluate and ones it can't +/// (e.g. lookaround/backreferences, which Rust's `regex` crate intentionally +/// doesn't support), then builds a single combined automaton from the former. +/// +/// Returns the prefilter plus the indices of `patterns` it does not cover; +/// callers must keep running those through the existing per-pattern path. +#[pyo3::pyfunction] +pub fn build_pattern_prefilter(patterns: Vec) -> (PatternPrefilter, Vec) { + let mut covered_sources: Vec<&str> = Vec::new(); + let mut covered_indices: Vec = Vec::new(); + let mut rejected_indices: Vec = Vec::new(); + + for (index, source) in patterns.iter().enumerate() { + match RegexBuilder::new(source).case_insensitive(true).build() { + Ok(_) => { + covered_sources.push(source); + covered_indices.push(index); + } + Err(_) => rejected_indices.push(index), + } + } + + if covered_sources.is_empty() { + return (PatternPrefilter { regex_set: None }, rejected_indices); + } + + match RegexSetBuilder::new(&covered_sources) + .case_insensitive(true) + .build() + { + Ok(set) => ( + PatternPrefilter { + regex_set: Some(set), + }, + rejected_indices, + ), + Err(_) => { + rejected_indices.extend(covered_indices); + (PatternPrefilter { regex_set: None }, rejected_indices) + } + } +}