mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat(rust): add litellm_core PyO3 extension with regex pattern prefilter
Signed-off-by: Jacinta Ferrant <jacinta.ferrant@gmail.com>
This commit is contained in:
Jacinta Ferrant
parent
78a7d0b210
commit
2d77cc368c
6 changed files with 281 additions and 0 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -1,5 +1,6 @@
|
|||
.python-version
|
||||
.venv
|
||||
rust/target/
|
||||
.venv_policy_test
|
||||
.env
|
||||
.claude
|
||||
|
|
|
|||
177
rust/Cargo.lock
generated
Normal file
177
rust/Cargo.lock
generated
Normal file
|
|
@ -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"
|
||||
12
rust/Cargo.toml
Normal file
12
rust/Cargo.toml
Normal file
|
|
@ -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"
|
||||
12
rust/pyproject.toml
Normal file
12
rust/pyproject.toml
Normal file
|
|
@ -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"]
|
||||
16
rust/src/lib.rs
Normal file
16
rust/src/lib.rs
Normal file
|
|
@ -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(())
|
||||
}
|
||||
63
rust/src/pattern_prefilter.rs
Normal file
63
rust/src/pattern_prefilter.rs
Normal file
|
|
@ -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<RegexSet>,
|
||||
}
|
||||
|
||||
#[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<String>) -> (PatternPrefilter, Vec<usize>) {
|
||||
let mut covered_sources: Vec<&str> = Vec::new();
|
||||
let mut covered_indices: Vec<usize> = Vec::new();
|
||||
let mut rejected_indices: Vec<usize> = 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue