diff --git a/Cargo.lock b/Cargo.lock index 70caa4bb..2f1f2a6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -179,12 +179,14 @@ dependencies = [ "base64 0.22.1", "chrono", "hmac 0.12.1", + "ort", "regex", "reqwest 0.12.28", "serde", "serde_json", "sha1", "thiserror 1.0.69", + "tokenizers", "tokio", "tracing", "tracing-subscriber", @@ -1216,6 +1218,12 @@ dependencies = [ "fastrand", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.21.7" @@ -1417,6 +1425,15 @@ dependencies = [ "either", ] +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.2.60" @@ -1551,6 +1568,21 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + [[package]] name = "compression-codecs" version = "0.4.37" @@ -1665,6 +1697,16 @@ dependencies = [ "crossbeam-utils", ] +[[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.18" @@ -1708,14 +1750,44 @@ dependencies = [ "cmov", ] +[[package]] +name = "daachorse" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f55d7153ba3b507595872a3874803f07a8a81d1e888abed8e5db7da0597d6e2" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + [[package]] name = "darling" version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", ] [[package]] @@ -1731,17 +1803,37 @@ dependencies = [ "syn", ] +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn", +] + [[package]] name = "darling_macro" version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ - "darling_core", + "darling_core 0.23.0", "quote", "syn", ] +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "dashmap" version = "6.1.0" @@ -1804,6 +1896,37 @@ dependencies = [ "serde_core", ] +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -1901,6 +2024,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + [[package]] name = "etcd-client" version = "0.14.1" @@ -2323,6 +2452,12 @@ dependencies = [ "sha1", ] +[[package]] +name = "hmac-sha256" +version = "1.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" + [[package]] name = "home" version = "0.5.12" @@ -2937,6 +3072,28 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "lzma-rust2" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e20f57f9918e5bd7bc58c22cdd70a6afc7375d4dd9683af5f2b34bd3d2bba619" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c" +dependencies = [ + "macro_rules_attribute-proc_macro", + "pastey", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" + [[package]] name = "matchers" version = "0.2.0" @@ -2958,6 +3115,16 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "md-5" version = "0.10.6" @@ -3090,6 +3257,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "multer" version = "3.1.0" @@ -3113,6 +3302,21 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "nom" version = "7.1.3" @@ -3296,6 +3500,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "openssl-probe" version = "0.1.6" @@ -3308,6 +3534,30 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "ort" +version = "2.0.0-rc.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4336a1e2b38848325241c72889086886004e589b7c74f335e60a8e8db5138a0b" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec", + "tracing", + "ureq", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf211e3776eea6aec988552fa118dd746d70e1b1e5e244058d1c98015f3e5872" +dependencies = [ + "hmac-sha256", + "lzma-rust2", + "ureq", +] + [[package]] name = "outref" version = "0.5.2" @@ -3462,12 +3712,27 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + [[package]] name = "portable-atomic" version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -3832,6 +4097,43 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[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-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[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 = "rcgen" version = "0.13.2" @@ -4497,7 +4799,7 @@ version = "3.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" dependencies = [ - "darling", + "darling 0.23.0", "proc-macro2", "quote", "syn", @@ -4654,12 +4956,35 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + [[package]] name = "spin" version = "0.9.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + [[package]] name = "sse-stream" version = "0.2.5" @@ -4679,6 +5004,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "string-builder" version = "0.2.0" @@ -4952,6 +5283,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44e5bea67576e04b6ff8564c5d9e09c2ef0cf476502245f2f120e497769d3112" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.4", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.52.1" @@ -5357,12 +5721,33 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + [[package]] name = "unicode-xid" version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "untrusted" version = "0.7.1" @@ -5386,6 +5771,7 @@ dependencies = [ "percent-encoding", "rustls", "rustls-pki-types", + "socks", "ureq-proto", "utf8-zero", "webpki-roots 1.0.7", diff --git a/Cargo.toml b/Cargo.toml index 777723e3..dc424d91 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -96,6 +96,18 @@ schemars = "0.8" # network access at runtime. tiktoken-rs = "0.12" +# Local-model guardrail (AISIX-Cloud#1331): in-process CPU embedding +# inference over ONNX Runtime. `tls-rustls` (not the default `tls-native`) +# keeps the build-time binary downloader off OpenSSL, matching the +# workspace's pure-rustls TLS stance. Offline builds point +# `ORT_LIB_LOCATION` at a pre-fetched ONNX Runtime; see +# aisix-guardrails/src/local_model.rs. +ort = { version = "=2.0.0-rc.13", default-features = false, features = ["std", "tracing", "download-binaries", "tls-rustls", "copy-dylibs", "api-27"] } +# HF tokenizer runtime for the model's tokenizer.json. `onig` is the regex +# engine its pre-tokenizer split patterns need; default `progressbar` / +# `esaxx_fast` are training/CLI conveniences an inference path never uses. +tokenizers = { version = "0.23", default-features = false, features = ["onig"] } + # Concurrency primitives arc-swap = "1.7" dashmap = "6.1" diff --git a/crates/aisix-guardrails/Cargo.toml b/crates/aisix-guardrails/Cargo.toml index 6ff7dba4..20cdc8a8 100644 --- a/crates/aisix-guardrails/Cargo.toml +++ b/crates/aisix-guardrails/Cargo.toml @@ -44,6 +44,13 @@ base64 = { workspace = true, optional = true } chrono = { workspace = true, optional = true } uuid = { workspace = true, optional = true } +# Local-model guardrail MVP (AISIX-Cloud#1331): in-process CPU embedding +# inference (ONNX Runtime + HF tokenizer). NOT in default features — it +# statically links ONNX Runtime (tens of MB) and is env-var-activated +# experimental surface; builds opt in with `--features local-model`. +ort = { workspace = true, optional = true } +tokenizers = { workspace = true, optional = true } + [features] default = [ "bedrock", @@ -77,6 +84,10 @@ aliyun-text-moderation = [ lakera = ["dep:reqwest"] openai-moderation = ["dep:reqwest"] presidio = ["dep:reqwest"] +# Local CPU embedding-model guardrail MVP (AISIX-Cloud#1331). Off by +# default: statically links ONNX Runtime and is activated only by env +# config (no control-plane surface yet). +local-model = ["dep:ort", "dep:tokenizers"] [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/crates/aisix-guardrails/src/lib.rs b/crates/aisix-guardrails/src/lib.rs index 8508946f..6d44b398 100644 --- a/crates/aisix-guardrails/src/lib.rs +++ b/crates/aisix-guardrails/src/lib.rs @@ -30,6 +30,8 @@ mod index; mod keyword; #[cfg(feature = "lakera")] mod lakera; +#[cfg(feature = "local-model")] +mod local_model; #[cfg(feature = "openai-moderation")] mod openai_moderation; mod pii; @@ -196,6 +198,8 @@ pub use index::{GuardrailIndex, RequestContext}; pub use keyword::{KeywordBlocklist, KeywordRule}; #[cfg(feature = "lakera")] pub use lakera::LakeraGuardrail; +#[cfg(feature = "local-model")] +pub use local_model::{LocalModelConfig, LocalModelError, LocalModelGuardrail, MODEL_DIR_ENV}; #[cfg(feature = "openai-moderation")] pub use openai_moderation::OpenaiModerationGuardrail; pub use pii::{builtin_rule, PiiAction, PiiGuardrail, PiiRule, BUILTIN_DETECTORS}; diff --git a/crates/aisix-guardrails/src/local_model.rs b/crates/aisix-guardrails/src/local_model.rs new file mode 100644 index 00000000..c62a7007 --- /dev/null +++ b/crates/aisix-guardrails/src/local_model.rs @@ -0,0 +1,842 @@ +//! Local CPU embedding-model guardrail — MVP vertical slice of the +//! second-tier guardrail (AISIX-Cloud#1331). +//! +//! Proves ONE thing: an in-process ONNX embedding model can sit inside +//! the guardrail chain, produce a span-level judgement, and drive a real +//! in-place rewrite. It is NOT the full three-layer pipeline from the +//! design issue (no rule scoring layer, no prototype library, no +//! standard risk categories, single hardcoded category). +//! +//! Pipeline per text segment: +//! 1. regex finds candidate spans with exact byte offsets +//! (dotted number runs, the EDA-version candidate shape); +//! 2. a context window around each candidate (±[`WINDOW_CONTEXT_CHARS`] +//! chars — the keyword-proximity window magnitude mainstream DLP +//! engines use, typically 50–300 chars) is embedded by the local +//! model and compared, by cosine similarity, against ONE category +//! prototype vector encoded at load time from a hardcoded Chinese +//! description sentence; +//! 3. above-threshold candidates are rewritten in place to +//! [`MASK_REPLACEMENT`]; everything else is returned byte-identical. +//! +//! The verdict is always `Allow` — this guardrail rewrites, never blocks +//! (the design issue's "只改写不阻断" hard constraint). +//! +//! Wire-in: implements the async segment-moderation hooks +//! ([`Guardrail::moderate_input_segments`] / +//! [`Guardrail::moderate_output_segments`]) — the same mask write-back +//! channel the Bedrock ANONYMIZE pass uses — so the proxy's existing +//! collect→moderate→apply walkers do the per-field rewrite and no proxy +//! plumbing changes. The sync `redact_*_text` hooks stay unimplemented: +//! inference must not run inline on a tokio worker. +//! +//! Threading (inherited from the #1271 assessment): inference runs in +//! `spawn_blocking`, bounded by a semaphore sized to the session pool — +//! the configured LANES (`GUARDRAIL_LOCAL_MODEL_LANES`, default 1; +//! api7/aisix#1001). Each session runs with ONE intra-op thread and +//! intra/inter-op **spinning disabled** — ONNX Runtime's default +//! spin-wait burns ~9.4% of a core while completely idle, taxing +//! deployments that never send guardrail traffic. The request's async +//! worker never blocks: it awaits the JoinHandle and keeps serving its +//! other connections; only blocking-pool threads compute. Those threads +//! are NOT core-pinned — on a saturated host they compete with the +//! serving workers for scheduling; hard business/model core +//! partitioning needs a dedicated pinned inference pool. +//! +//! Scaling notes (MVP review, measured on a 12-core avx2+vnni host): +//! - One lane sustains ~50 inferences/s (p50 ≈ 19 ms per ~35-token +//! window; roughly linear in tokens). The acceptance shape spends TWO +//! inferences per request (input + output pass). +//! - Throughput scales by adding lanes (api7/aisix#1001, implemented): +//! `GUARDRAIL_LOCAL_MODEL_LANES` = N sessions behind +//! `Semaphore::new(N)`. `run(&mut self)` forbids concurrent runs on +//! one session even though the ONNX Runtime C API documents `Run` as +//! thread-safe with shared read-only weights; this crate forbids +//! `unsafe`, so the shared-weights form waits on an upstream `&self` +//! run, a thin unsafe shim crate, or the sidecar deployment form. +//! Until then each lane pays its own weight copy — measured: the +//! first lane costs ~192 MiB resident (weights + tokenizer + arena), +//! each additional lane ~102 MiB (its weight copy + arena). +//! Lane dispatch is a centralized free-list — deliberately NO +//! worker↔session binding (sessions are interchangeable; the central +//! queue load-balances the uneven per-worker accept distribution). +//! ONE ORT `Session` is a loaded model instance, not a conversation: +//! every `run` is stateless, so lanes are freely interchangeable. +//! - Per-audit cost ≈ (#candidate windows × inference) + (#prototypes × +//! 384-dim dot ≈ 1 µs). A window embedding is category-agnostic: a +//! span matched by several categories embeds ONCE and compares +//! against the whole prototype library, so prototype count stays +//! latency-noise up to ~1e5 vectors (then: ANN index). +//! - The candidate-regex layer is the all-traffic cost. At many +//! categories the per-category patterns must merge into one +//! multi-pattern automaton (`RegexSet` / aho-corasick), compiled at +//! prototype-library build time and atomically swapped — never per +//! request. The rust regex engine is non-backtracking, so +//! operator-supplied patterns cannot ReDoS the data plane. +//! +//! Model contract (upstream docs): the model directory holds the two +//! deliverables `model.onnx` + `tokenizer.json`. The MVP target is +//! `ibm-granite/granite-embedding-97m-multilingual-r2`'s official int8 +//! ONNX export (`onnx/model_quint8_avx2.onnx`, 93.7 MiB, standard +//! `ai.onnx` opset only — the q4/q4f16/bnb4 variants carry +//! `com.microsoft` ops and are ruled out). Per the repo's +//! `1_Pooling/config.json` + `modules.json`, sentence embedding = CLS +//! pooling over `last_hidden_state` followed by L2 normalization; the +//! ModernBERT graph takes `input_ids` + `attention_mask` (no +//! `token_type_ids`). +//! +//! +//! Offline builds: `ort-sys` honors `ORT_OFFLINE=1` (skip the prebuilt +//! ONNX Runtime download) with `ORT_LIB_PATH` pointing at a pre-fetched +//! library (see `ort-sys` `build/vars.rs`). + +use std::collections::BTreeMap; +use std::ops::Range; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use async_trait::async_trait; +use ort::session::Session; +use ort::value::Tensor; +use regex::Regex; + +use crate::{ + Guardrail, GuardrailVerdict, SegmentsOutcome, StreamOutputPolicy, + DEFAULT_STREAM_OUTPUT_BUFFER_BYTES, +}; + +/// Environment variable holding the model directory (`model.onnx` + +/// `tokenizer.json`). Set → the server bootstrap loads and injects the +/// guardrail; unset → the feature is completely inert. +pub const MODEL_DIR_ENV: &str = "GUARDRAIL_LOCAL_MODEL_DIR"; +/// Optional cosine-similarity gate override (default +/// [`DEFAULT_THRESHOLD`]). +pub const THRESHOLD_ENV: &str = "GUARDRAIL_LOCAL_MODEL_THRESHOLD"; +/// Optional inference-lane count (default 1, clamped to +/// [`MAX_LANES`]). Each lane is one ONNX session — one more core the +/// guardrail may use and roughly one more ~100 MiB weight copy resident (measured); see +/// the module scaling notes and api7/aisix#1001. +pub const LANES_ENV: &str = "GUARDRAIL_LOCAL_MODEL_LANES"; + +/// Upper clamp for [`LANES_ENV`]: lanes are cores, and no sane host +/// grants the guardrail more than this. +const MAX_LANES: usize = 32; + +/// Cosine-similarity gate for "this window is about the category". +/// Calibrated with this module's `#[ignore]` probe against the +/// prototype below: the acceptance-style positive scores ~0.90, every +/// probed negative (compile-log timings, memory sizes, IPs, plain +/// numbers) ≤0.76. Deliberately precision-leaning — a mask +/// false-positive corrupts user content — at a measured recall cost: +/// harder positives ("升级到 2022.4"、"Virtuoso IC6.1.8") score +/// 0.75–0.79 and are NOT masked at this default. Single-prototype +/// zero-shot cosine cannot separate those from the hard negatives at +/// all (every phrasing swept had negative margin); closing that gap is +/// the design issue's rule-scoring layer + real-sample prototypes, not +/// a threshold tweak. +const DEFAULT_THRESHOLD: f32 = 0.80; + +/// The ONE category this MVP ships: EDA-software version numbers. The +/// prototype vector is the load-time embedding of this sentence — the +/// v1 "customer types one Chinese description" path from the design +/// issue, collapsed to a compile-time constant. +const PROTOTYPE_DESCRIPTION_ZH: &str = "EDA 软件的版本号"; + +/// Candidate shape: a dotted number run (`12.1`, `2022.4`, `6.1.8`). +/// Plain integers are out of MVP scope. +const CANDIDATE_PATTERN: &str = r"\d+(?:\.\d+)+"; + +/// Context chars kept on each side of a candidate when cutting the +/// window the model judges. +const WINDOW_CONTEXT_CHARS: usize = 50; + +/// Hard cap on model invocations per moderation pass (the design +/// issue's "每请求模型调用次数上限" recommendation). Candidates past +/// the cap are left untouched and a warning is logged — degrade to +/// doing less, never to blocking or stalling. +const MAX_MODEL_CALLS_PER_PASS: usize = 8; + +/// Hardcoded rewrite for an above-threshold candidate span. +const MASK_REPLACEMENT: &str = "***"; + +/// Detector label used in redaction counts (never the matched value — +/// #153 / #932 no-leak rule). +const DETECTOR_NAME: &str = "eda_version"; + +#[derive(Debug, thiserror::Error)] +pub enum LocalModelError { + #[error("local-model guardrail: tokenizer: {0}")] + Tokenizer(String), + #[error("local-model guardrail: onnx runtime: {0}")] + Ort(#[from] ort::Error), + #[error("local-model guardrail: {0}")] + Model(String), +} + +/// Load-time configuration. MVP surface: env vars only — deliberately +/// NO control-plane resource; see the design issue for what the real +/// config object must eventually carry (model id + dimension + prefix +/// convention + threshold table + prototype-library version). +/// +/// Target form (not built): the category data becomes a prototype +/// LIBRARY resource — etcd stores category TEXT only (name, action, +/// replacement, candidate patterns, description/sample sentences); the +/// DP encodes vectors with its own loaded model off the hot path +/// (cached by model id + text hash) and atomically swaps the compiled +/// library on a watch tick, the same propagation the guardrail index +/// already uses — so category updates land without a restart. Vectors +/// never travel through config, which removes the "encoded by which +/// model" mismatch for text-sourced prototypes; swapping the MODEL +/// itself stays a cold operation (all vectors rebuilt). +pub struct LocalModelConfig { + pub model_dir: PathBuf, + pub threshold: f32, + /// Inference lanes = ONNX sessions = max concurrent inferences. + pub lanes: usize, +} + +impl LocalModelConfig { + /// `None` when [`MODEL_DIR_ENV`] is unset (guardrail disabled). A + /// malformed or out-of-range threshold falls back to the default + /// rather than failing boot — the gate is a tuning knob, not a + /// correctness one. The range check matters: `"NaN"` parses as a + /// valid f32 and would make `score >= threshold` always false — a + /// configured-looking guardrail that silently never masks. Lanes + /// follow the same lenient rule (malformed → 1). + pub fn from_env() -> Option { + let model_dir = PathBuf::from(std::env::var_os(MODEL_DIR_ENV)?); + let threshold = std::env::var(THRESHOLD_ENV) + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|t| t.is_finite() && (0.0..=1.0).contains(t)) + .unwrap_or(DEFAULT_THRESHOLD); + let lanes = parse_lanes(std::env::var(LANES_ENV).ok().as_deref()); + Some(Self { + model_dir, + threshold, + lanes, + }) + } +} + +/// [`LANES_ENV`] parse rule: default 1 when unset or malformed (zero +/// included — a zero-lane guardrail is a misconfiguration, not a +/// disable switch; disabling is unsetting [`MODEL_DIR_ENV`]), clamped +/// to [`MAX_LANES`]. +fn parse_lanes(raw: Option<&str>) -> usize { + raw.and_then(|s| s.parse::().ok()) + .filter(|&n| n >= 1) + .unwrap_or(1) + .min(MAX_LANES) +} + +/// The blocking inference core: one shared tokenizer (`encode` takes +/// `&self` and is thread-safe) + a pool of ONNX sessions — the +/// inference LANES (api7/aisix#1001). Each session sits behind its own +/// mutex (`Session::run` takes `&mut`); the guardrail's semaphore +/// admits at most `sessions.len()` concurrent inferences, so an +/// admitted task always finds a free session. Dispatch is a +/// centralized free-list: any request takes whichever session is idle — +/// deliberately NO worker↔session binding (sessions are stateless and +/// interchangeable; a central queue load-balances the uneven per-worker +/// accept distribution). +struct Embedder { + tokenizer: tokenizers::Tokenizer, + sessions: Vec>, +} + +impl Embedder { + fn load(dir: &std::path::Path, lanes: usize) -> Result { + // `parse_lanes` never yields 0, but `LocalModelConfig`'s fields + // are `pub` and a future programmatic constructor (the + // control-plane integration) could hand-build one; an empty + // pool would panic at the prototype embed with a confusing + // boot error instead of just working. + let lanes = lanes.max(1); + let tokenizer_path = dir.join("tokenizer.json"); + let mut tokenizer = tokenizers::Tokenizer::from_file(&tokenizer_path).map_err(|e| { + LocalModelError::Tokenizer(format!("{}: {e}", tokenizer_path.display())) + })?; + // Enforce truncation here instead of trusting the operator's + // tokenizer.json (the reference export allows 32K tokens): a + // window is ~100 chars by construction, so 512 tokens + // (`TruncationParams` default) is pure headroom — this is the + // second bound, after the candidate-span byte cap, that keeps a + // single inference's cost fixed no matter what the caller sends. + tokenizer + .with_truncation(Some(tokenizers::TruncationParams::default())) + .map_err(|e| LocalModelError::Tokenizer(format!("truncation config: {e}")))?; + let model_path = dir.join("model.onnx"); + // One intra-op thread per session and NO spin-wait: throughput + // scales by LANES (each lane single-threaded is the best + // per-core efficiency for short windows), and the guardrail + // must not tax a gateway that isn't sending it traffic (see + // module doc). Builder-stage errors carry the builder back for + // recovery (`ort::Error`); this path never + // recovers, so they fold to their message. + let build = |b: ort::session::builder::SessionBuilder| { + b.with_intra_threads(1)? + .with_inter_threads(1)? + .with_intra_op_spinning(false)? + .with_inter_op_spinning(false) + }; + let mut sessions = Vec::with_capacity(lanes); + for _ in 0..lanes { + let session = build(Session::builder()?) + .map_err(|e| LocalModelError::Model(format!("session options: {e}")))? + .commit_from_file(&model_path)?; + sessions.push(Mutex::new(session)); + } + Ok(Self { + tokenizer, + sessions, + }) + } + + /// Take an idle session from the pool. The caller holds a semaphore + /// permit and permits == sessions, so a free lane exists whenever + /// this runs; the final blocking `lock()` is a defensive fallback + /// that cannot deadlock (some lane always releases). A panic in a + /// previous run poisons that lane's mutex; recover instead of + /// panicking forever after — `Session::run` is stateless, so the + /// session itself is unharmed, and turning every later inference + /// into a panic would silently disable the operator's masking until + /// restart (the exact state load treats as boot-fatal). + fn acquire_free_session(&self) -> std::sync::MutexGuard<'_, Session> { + for lane in &self.sessions { + match lane.try_lock() { + Ok(guard) => return guard, + Err(std::sync::TryLockError::Poisoned(poisoned)) => { + return poisoned.into_inner(); + } + Err(std::sync::TryLockError::WouldBlock) => continue, + } + } + self.sessions[0] + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Embed one text: encode (with the tokenizer's special-token + /// template), run the graph, CLS-pool `last_hidden_state`, L2 + /// normalize. Blocking — call from `spawn_blocking` only. + fn embed(&self, text: &str) -> Result, LocalModelError> { + let encoding = self + .tokenizer + .encode(text, true) + .map_err(|e| LocalModelError::Tokenizer(e.to_string()))?; + let ids: Vec = encoding.get_ids().iter().map(|&t| i64::from(t)).collect(); + let mask: Vec = encoding + .get_attention_mask() + .iter() + .map(|&t| i64::from(t)) + .collect(); + let len = ids.len() as i64; + + let started = Instant::now(); + let mut session = self.acquire_free_session(); + let outputs = session.run(ort::inputs! { + "input_ids" => Tensor::from_array((vec![1, len], ids))?, + "attention_mask" => Tensor::from_array((vec![1, len], mask))?, + })?; + let (shape, data) = outputs + .get("last_hidden_state") + .ok_or_else(|| LocalModelError::Model("model has no last_hidden_state output".into()))? + .try_extract_tensor::()?; + // Expected [1, seq, hidden]; CLS pooling = the first hidden-size + // slice (the template's leading special token). + if shape.len() != 3 { + return Err(LocalModelError::Model(format!( + "last_hidden_state has rank {} (expected 3)", + shape.len() + ))); + } + let hidden = shape[2] as usize; + let mut cls = data + .get(..hidden) + .ok_or_else(|| LocalModelError::Model("empty last_hidden_state".into()))? + .to_vec(); + let norm = cls.iter().map(|v| v * v).sum::().sqrt(); + if norm > 0.0 { + for v in &mut cls { + *v /= norm; + } + } + tracing::debug!( + infer_us = started.elapsed().as_micros() as u64, + tokens = encoding.get_ids().len(), + "local-model guardrail inference" + ); + Ok(cls) + } +} + +/// Dot product of two L2-normalized vectors = cosine similarity. +fn cosine(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() +} + +/// Hard byte cap on a single candidate span. A real version number is +/// short by definition; without this cap `\d+(?:\.\d+)+` matches an +/// arbitrarily long `1.1.1...` run as ONE span, and since the window is +/// span + context, a crafted request would turn each model call into a +/// max-truncation inference and stall the single lane for everyone +/// (audit finding on PR #999). Over-cap spans are dropped BEFORE the +/// per-pass budget so they cannot starve legitimate candidates either. +const MAX_CANDIDATE_SPAN_BYTES: usize = 64; + +/// Candidate spans (byte ranges) in `text`, in order. Spans longer than +/// [`MAX_CANDIDATE_SPAN_BYTES`] are not candidates (see the constant). +fn candidate_spans(re: &Regex, text: &str) -> Vec> { + re.find_iter(text) + .map(|m| m.range()) + .filter(|s| s.len() <= MAX_CANDIDATE_SPAN_BYTES) + .collect() +} + +/// The context window around `span`: `ctx` chars on each side, snapped +/// to char boundaries, clamped to the text. +fn window_bounds(text: &str, span: &Range, ctx: usize) -> Range { + if ctx == 0 { + return span.clone(); + } + let before = &text[..span.start]; + let start = before + .char_indices() + .rev() + .nth(ctx - 1) + .map_or(0, |(i, _)| i); + let after = &text[span.end..]; + let end = span.end + + after + .char_indices() + .nth(ctx) + .map_or(after.len(), |(i, _)| i); + start..end +} + +/// Rewrite `spans` (ascending, non-overlapping — regex `find_iter` +/// order) in `text` to [`MASK_REPLACEMENT`], right-to-left so earlier +/// offsets stay valid. +fn apply_masks(text: &str, spans: &[Range]) -> String { + let mut out = text.to_owned(); + for span in spans.iter().rev() { + out.replace_range(span.clone(), MASK_REPLACEMENT); + } + out +} + +/// The runtime guardrail. Always-`Allow`; masks via the segment hooks. +pub struct LocalModelGuardrail { + embedder: Arc, + /// L2-normalized embedding of [`PROTOTYPE_DESCRIPTION_ZH`]. + prototype: Vec, + threshold: f32, + candidate_re: Regex, + /// Bounds in-flight `spawn_blocking` inference tasks. Sized to the + /// session-pool size (the configured lanes): more permits would + /// only queue on the session mutexes from inside blocking threads. + permits: Arc, +} + +impl LocalModelGuardrail { + /// Load tokenizer + the session pool and encode the category + /// prototype. Blocking (N model loads + one inference) — the server + /// bootstrap wraps it in `spawn_blocking`. + pub fn load(config: &LocalModelConfig) -> Result { + let started = Instant::now(); + let embedder = Embedder::load(&config.model_dir, config.lanes)?; + let prototype = embedder.embed(PROTOTYPE_DESCRIPTION_ZH)?; + tracing::info!( + model_dir = %config.model_dir.display(), + threshold = config.threshold, + lanes = config.lanes, + load_ms = started.elapsed().as_millis() as u64, + "local-model guardrail loaded (category: EDA software version)" + ); + Ok(Self { + embedder: Arc::new(embedder), + prototype, + threshold: config.threshold, + candidate_re: Regex::new(CANDIDATE_PATTERN).expect("candidate pattern must compile"), + permits: Arc::new(tokio::sync::Semaphore::new(config.lanes)), + }) + } + + /// Embed one window off the async runtime: bounded by the + /// semaphore, executed in `spawn_blocking`. The permit MOVES INTO + /// the blocking closure: a `spawn_blocking` task keeps running when + /// its awaiter is dropped (client disconnect), so a permit held in + /// this async scope would release while the session is still busy — + /// letting the next request park a second blocking-pool thread on + /// the session mutex, and repeated cancellations pile threads up + /// toward the pool cap (audit finding on PR #999). Held by the + /// closure, the permit releases only when the inference actually + /// finishes. + async fn embed_window(&self, window: String) -> Result, LocalModelError> { + let permit = Arc::clone(&self.permits) + .acquire_owned() + .await + .expect("inference semaphore never closed"); + let embedder = Arc::clone(&self.embedder); + tokio::task::spawn_blocking(move || { + let _permit = permit; + embedder.embed(&window) + }) + .await + .map_err(|e| LocalModelError::Model(format!("inference task join: {e}")))? + } + + /// Mask one segment. Returns the rewritten text and how many spans + /// were masked; `budget` is the shared per-pass model-call cap. + /// Model failure on a candidate leaves that span untouched (rewrite + /// less, never block — the fail-open arm of "只改写不阻断"). + async fn mask_segment(&self, text: &str, budget: &mut usize) -> (String, u32) { + let mut hits: Vec> = Vec::new(); + for span in candidate_spans(&self.candidate_re, text) { + if *budget == 0 { + tracing::warn!( + cap = MAX_MODEL_CALLS_PER_PASS, + "local-model guardrail: candidate cap reached; remaining candidates left unmasked" + ); + break; + } + *budget -= 1; + let window = text[window_bounds(text, &span, WINDOW_CONTEXT_CHARS)].to_owned(); + match self.embed_window(window).await { + Ok(vector) => { + let score = cosine(&vector, &self.prototype); + tracing::debug!( + score, + threshold = self.threshold, + "local-model window judged" + ); + if score >= self.threshold { + hits.push(span); + } + } + Err(err) => { + tracing::warn!(error = %err, "local-model guardrail inference failed; span left unmasked"); + } + } + } + if hits.is_empty() { + (text.to_owned(), 0) + } else { + let count = hits.len() as u32; + (apply_masks(text, &hits), count) + } + } + + /// One moderation pass over a request's (or response's) text + /// segments — the shared body of both segment hooks. + async fn moderate(&self, texts: &[String]) -> SegmentsOutcome { + let mut budget = MAX_MODEL_CALLS_PER_PASS; + let mut masked: Vec = Vec::with_capacity(texts.len()); + let mut total = 0u32; + for text in texts { + let (rewritten, count) = self.mask_segment(text, &mut budget).await; + masked.push(rewritten); + total += count; + } + let mut counts = BTreeMap::new(); + if total > 0 { + counts.insert(DETECTOR_NAME.to_owned(), total); + } + SegmentsOutcome { + verdict: GuardrailVerdict::Allow, + masked: (total > 0).then_some(masked), + counts, + monitor_hits: Vec::new(), + } + } +} + +#[async_trait] +impl Guardrail for LocalModelGuardrail { + fn name(&self) -> &'static str { + "local_model" + } + + /// Consulted through the segment hooks only (the mask write-back + /// channel); the plain `check_*` hooks stay default-`Allow`. + fn moderates_segments(&self) -> bool { + true + } + + /// Masking a streamed response needs the whole response held back (a + /// span can cross any chunk boundary) — but past the buffer cap this + /// guardrail must release UNMASKED, not block: the trait default's + /// fail-closed overflow would turn a >cap streamed response into a + /// `content_filter` error from a guardrail whose contract is + /// "rewrite, never block" (audit finding on PR #999). Past-cap + /// content degrades to fewer masks, the same fail-open arm as + /// inference failure and the per-pass call cap. A chain member with + /// a stricter policy (e.g. fail-closed pii) still wins the fold. + fn stream_output_policy(&self) -> StreamOutputPolicy { + StreamOutputPolicy::BufferFull { + max_buffer_bytes: DEFAULT_STREAM_OUTPUT_BUFFER_BYTES, + on_exceeded_fail_open: true, + } + } + + async fn moderate_input_segments(&self, texts: &[String]) -> SegmentsOutcome { + self.moderate(texts).await + } + + async fn moderate_output_segments(&self, texts: &[String]) -> SegmentsOutcome { + self.moderate(texts).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn re() -> Regex { + Regex::new(CANDIDATE_PATTERN).unwrap() + } + + #[test] + fn candidate_spans_find_dotted_runs_only() { + let text = "版本是 12.1,构建号 2022.4.1,端口 8080"; + let spans = candidate_spans(&re(), text); + let values: Vec<&str> = spans.iter().map(|s| &text[s.clone()]).collect(); + assert_eq!(values, vec!["12.1", "2022.4.1"]); + } + + #[test] + fn candidate_spans_drop_oversized_runs() { + // A crafted `1.1.1...` run over the byte cap is one regex match + // but NOT a candidate — it must vanish before budget accounting + // so it can neither stall the lane nor starve real candidates. + let bomb = "1.1".repeat(60); // 180 bytes, single match + let text = format!("前缀 {bomb} 中缀 12.1 后缀"); + let spans = candidate_spans(&re(), &text); + let values: Vec<&str> = spans.iter().map(|s| &text[s.clone()]).collect(); + assert_eq!(values, vec!["12.1"]); + } + + #[test] + fn window_bounds_snap_to_char_boundaries() { + let text = "这个 EDA 软件的版本是 12.1,请勿外传"; + let span = candidate_spans(&re(), text).remove(0); + // A tiny context still lands on char boundaries around CJK. + let w = window_bounds(text, &span, 3); + let window = &text[w]; + assert!(window.contains("12.1"), "window: {window}"); + assert_eq!(window, "本是 12.1,请勿"); + } + + #[test] + fn window_bounds_clamp_to_text_edges() { + let text = "12.1 只有后文"; + let span = candidate_spans(&re(), text).remove(0); + let w = window_bounds(text, &span, 50); + assert_eq!(&text[w], text); + } + + #[test] + fn apply_masks_rewrites_right_to_left() { + let text = "从 12.1 升到 13.0 了"; + let spans = candidate_spans(&re(), text); + assert_eq!(apply_masks(text, &spans), "从 *** 升到 *** 了"); + } + + #[test] + fn config_from_env_requires_model_dir() { + // Isolated var names are process-global; this test only checks the + // parse fallback path via the public constructor contract. + let cfg = LocalModelConfig { + model_dir: PathBuf::from("/nonexistent"), + threshold: DEFAULT_THRESHOLD, + lanes: 1, + }; + assert!(LocalModelGuardrail::load(&cfg).is_err()); + } + + #[test] + fn parse_lanes_defaults_and_clamps() { + // Unset / malformed / zero → 1 (zero is a misconfiguration, not + // a disable switch); valid values pass; huge values clamp. + assert_eq!(parse_lanes(None), 1); + assert_eq!(parse_lanes(Some("")), 1); + assert_eq!(parse_lanes(Some("abc")), 1); + assert_eq!(parse_lanes(Some("-2")), 1); + assert_eq!(parse_lanes(Some("0")), 1); + assert_eq!(parse_lanes(Some("4")), 4); + assert_eq!(parse_lanes(Some("9999")), MAX_LANES); + } + + // ── model-backed tests (need the real model files) ────────────────── + // + // Run explicitly with the model directory present: + // GUARDRAIL_LOCAL_MODEL_DIR=~/.cache/aisix-local-guardrail-mvp \ + // cargo test -p aisix-guardrails --features local-model -- --ignored + + fn load_from_env() -> Option { + let cfg = LocalModelConfig::from_env()?; + Some(LocalModelGuardrail::load(&cfg).expect("model files present but load failed")) + } + + /// Calibration probe: prints the cosine matrix behind + /// [`DEFAULT_THRESHOLD`] and pins the MVP contract — the + /// acceptance-style positive clears the gate, every probed negative + /// stays under it. The harder positives are printed but NOT + /// asserted: single-prototype zero-shot cosine has no margin over + /// the hard negatives (measured; see the threshold doc), so at the + /// precision-leaning default they are a known recall gap. + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn probe_similarity_matrix() { + let Some(g) = load_from_env() else { return }; + let acceptance = "这个 EDA 软件的版本是 12.1,请确认兼容性"; + let recall_gap_positives = [ + "我们把仿真工具升级到 2022.4 之后速度快了很多", + "Virtuoso IC6.1.8 出现了崩溃", + ]; + let negatives = [ + "Elapsed: 12.345s, Memory: 4.2 GB", + "服务器的 IP 地址是 10.2.255.1", + "圆周率约等于 3.14159", + "工艺节点是 0.13um,良率还行", + ]; + let acc = cosine( + &g.embed_window(acceptance.to_owned()).await.unwrap(), + &g.prototype, + ); + println!("ACC {acc:.4} {acceptance}"); + for text in recall_gap_positives { + let s = cosine( + &g.embed_window(text.to_owned()).await.unwrap(), + &g.prototype, + ); + println!("POS(gap) {s:.4} {text}"); + } + assert!( + acc >= g.threshold, + "acceptance positive {acc:.4} under threshold {}", + g.threshold + ); + for text in negatives { + let s = cosine( + &g.embed_window(text.to_owned()).await.unwrap(), + &g.prototype, + ); + println!("NEG {s:.4} {text}"); + assert!( + s < g.threshold, + "negative {s:.4} would mask at threshold {}: {text}", + g.threshold + ); + } + } + + /// The acceptance path in miniature: the segment hook masks the + /// version number and only the version number. + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn masks_acceptance_sentence() { + let Some(g) = load_from_env() else { return }; + let texts = vec!["这个 EDA 软件的版本是 12.1".to_owned()]; + let outcome = g.moderate_input_segments(&texts).await; + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); + let masked = outcome.masked.expect("must mask the version number"); + assert_eq!(masked[0], "这个 EDA 软件的版本是 ***"); + assert_eq!(outcome.counts.get(DETECTOR_NAME), Some(&1)); + } + + /// Lanes are interchangeable and safe under concurrency: with a + /// 2-lane pool, 8 concurrent embeds of the same window all succeed + /// and agree (sessions share nothing but identical weights, and a + /// single-threaded run is deterministic). + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn lanes_run_concurrently_and_agree() { + let Some(mut cfg) = LocalModelConfig::from_env() else { + return; + }; + cfg.lanes = 2; + let g = std::sync::Arc::new( + LocalModelGuardrail::load(&cfg).expect("model files present but load failed"), + ); + let window = "这个 EDA 软件的版本是 12.1,请确认兼容性"; + let tasks: Vec<_> = (0..8) + .map(|_| { + let g = std::sync::Arc::clone(&g); + tokio::spawn(async move { g.embed_window(window.to_owned()).await }) + }) + .collect(); + let mut vectors = Vec::new(); + for t in tasks { + vectors.push(t.await.unwrap().expect("concurrent embed failed")); + } + let reference = &vectors[0]; + for v in &vectors[1..] { + let agreement = cosine(v, reference); + assert!(agreement > 0.9999, "lanes disagree: cosine {agreement}"); + } + } + + /// Throughput calibration for api7/aisix#1001: lanes come from + /// GUARDRAIL_LOCAL_MODEL_LANES, so one binary sweeps 1/2/4 lanes. + /// Prints inferences/s over a saturating concurrent batch. + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn probe_lane_throughput() { + let Some(g) = load_from_env().map(std::sync::Arc::new) else { + return; + }; + let window = "这个 EDA 软件的版本是 12.1,请确认与工艺库的兼容性之后再安排回归测试"; + // Warm-up. + for _ in 0..3 { + g.embed_window(window.to_owned()).await.unwrap(); + } + let total = 64usize; + let started = Instant::now(); + let tasks: Vec<_> = (0..total) + .map(|_| { + let g = std::sync::Arc::clone(&g); + tokio::spawn(async move { g.embed_window(window.to_owned()).await }) + }) + .collect(); + for t in tasks { + t.await.unwrap().expect("embed failed"); + } + let secs = started.elapsed().as_secs_f64(); + println!( + "lanes={} total={} wall={:.2}s throughput={:.1} inferences/s", + g.embedder.sessions.len(), + total, + secs, + total as f64 / secs + ); + } + + /// Rough single-inference latency figure for the MVP report. + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn probe_inference_latency() { + let Some(g) = load_from_env() else { return }; + let window = + "这个 EDA 软件的版本是 12.1,请确认与工艺库的兼容性之后再安排回归测试".to_owned(); + // Warm-up, then timed runs. + for _ in 0..3 { + g.embed_window(window.clone()).await.unwrap(); + } + let mut samples = Vec::new(); + for _ in 0..20 { + let t = Instant::now(); + g.embed_window(window.clone()).await.unwrap(); + samples.push(t.elapsed().as_micros() as u64); + } + samples.sort_unstable(); + println!( + "inference us over 20 runs: p50={} min={} max={}", + samples[samples.len() / 2], + samples[0], + samples[samples.len() - 1] + ); + } +} diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 3360f97b..fed90ffa 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1248,6 +1248,20 @@ async fn dispatch( *applied_out = applied_guardrails.clone(); let resolved_chain: std::sync::Arc = std::sync::Arc::new(resolved); + // AISIX-Cloud#1331 MVP: the env-injected local-model guardrail joins + // AFTER the attachment-resolved chain (its segment masks compose on + // the chain's output; nested-chain folds filter their own members, so + // wrapping is safe). Deployment-wide + mask-only, not a row: it has no + // `applied()` entry and never blocks. MVP wiring is chat-only; the + // sibling endpoint families are a tracked gap on the design issue. + let resolved_chain: std::sync::Arc = + match state.local_model_guardrail.as_ref() { + Some(local) => std::sync::Arc::new(aisix_guardrails::GuardrailChain::new(vec![ + std::sync::Arc::clone(&resolved_chain), + std::sync::Arc::clone(local), + ])), + None => resolved_chain, + }; // Input guardrails. Run before reservation so a blocked prompt // doesn't burn an RPM slot — content-policy refusals shouldn't diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index fe209afa..bc2cd737 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -190,6 +190,13 @@ pub struct ProxyStateInner { pub health: Arc, /// Public liveness state served on `GET /livez`. pub livez: Arc, + /// Env-injected local CPU embedding-model guardrail (AISIX-Cloud#1331 + /// MVP vertical slice). `None` (the default) = inactive. When present, + /// the chat handler composes it AFTER the per-request resolved chain: + /// it is deployment-wide experimental surface, not an + /// attachment-scoped guardrail row, and it only masks (never blocks). + /// MVP wiring covers `/v1/chat/completions` only. + pub local_model_guardrail: Option>, /// Runtime model-status tracker keyed by resolved direct-model id. /// Used for request-path cooldown/background health exclusion and /// surfaced by `GET /admin/v1/models/status`. @@ -294,6 +301,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::new()), livez: Arc::new(LivezState::new()), + local_model_guardrail: None, config_apply_age: None, runtime_status: Arc::new(ModelRuntimeStatusTracker::new()), usage_sink: UsageSink::disabled(), @@ -336,6 +344,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::new()), livez: Arc::new(LivezState::new()), + local_model_guardrail: None, config_apply_age: None, runtime_status: Arc::new(ModelRuntimeStatusTracker::new()), usage_sink: UsageSink::disabled(), @@ -392,6 +401,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::with_flags(bookkeeping_flags)), livez: Arc::new(LivezState::new()), + local_model_guardrail: None, config_apply_age: None, runtime_status, usage_sink: UsageSink::disabled(), @@ -432,6 +442,20 @@ impl ProxyState { self } + /// Inject the env-configured local-model guardrail (AISIX-Cloud#1331 + /// MVP). Wired by the server bootstrap when + /// `GUARDRAIL_LOCAL_MODEL_DIR` is set on a binary built with the + /// `local-model-guardrail` feature (non-`AISIX_` prefix on purpose: + /// the config loader maps every `AISIX_*` env var onto a config + /// field and strictly rejects unknown ones). + pub fn with_local_model_guardrail( + mut self, + guardrail: Arc, + ) -> Self { + Arc::make_mut(&mut self.inner).local_model_guardrail = Some(guardrail); + self + } + /// Swap in the classifier compiled from /// `observability.metrics.client_type_rules` (AISIX-Cloud#1045). /// Default is built-ins only. diff --git a/crates/aisix-server/Cargo.toml b/crates/aisix-server/Cargo.toml index c1f10950..87e483ba 100644 --- a/crates/aisix-server/Cargo.toml +++ b/crates/aisix-server/Cargo.toml @@ -68,6 +68,14 @@ hyper-util = { version = "0.1", features = ["server-auto", "tokio"] } tikv-jemallocator = "0.6" tikv-jemalloc-ctl = "0.6" +[features] +# Local CPU embedding-model guardrail MVP (AISIX-Cloud#1331). Off by +# default (statically links ONNX Runtime); the guardrail additionally +# activates only when GUARDRAIL_LOCAL_MODEL_DIR is set (non-AISIX_ +# prefix on purpose: the config loader claims the AISIX_* env namespace +# and strictly rejects unknown fields). +local-model-guardrail = ["aisix-guardrails/local-model"] + [dev-dependencies] tempfile = "3" wiremock = "0.6" diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 5dc08c30..f1eace6e 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -923,6 +923,28 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { bedrock_endpoint_url, Some(guardrail_metrics_sink), )); + // Local CPU embedding-model guardrail MVP (AISIX-Cloud#1331): + // env-activated, no control-plane surface yet. Load failure with the + // env var set is boot-fatal — a masking guardrail the operator asked + // for that silently isn't there would leak the very content it exists + // to rewrite. Model load + prototype inference block, so they run off + // the async bootstrap thread. + #[cfg(feature = "local-model-guardrail")] + if let Some(local_cfg) = aisix_guardrails::LocalModelConfig::from_env() { + let guardrail = tokio::task::spawn_blocking(move || { + aisix_guardrails::LocalModelGuardrail::load(&local_cfg) + }) + .await + .map_err(|e| anyhow::anyhow!("local-model guardrail load task: {e}"))??; + proxy_state = proxy_state.with_local_model_guardrail(Arc::new(guardrail)); + } + #[cfg(not(feature = "local-model-guardrail"))] + if std::env::var_os("GUARDRAIL_LOCAL_MODEL_DIR").is_some() { + tracing::warn!( + "GUARDRAIL_LOCAL_MODEL_DIR is set, but this binary was built without the \ + `local-model-guardrail` feature; ignoring" + ); + } // Heartbeat worker — spawned after proxy_state exists so it can read // the exporter fan-out's delivery counters. Each tick reports: // - rejected_resources: the supervisor's loader rejections (#115) diff --git a/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts b/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts new file mode 100644 index 00000000..6ae0ef43 --- /dev/null +++ b/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts @@ -0,0 +1,152 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + ProxyClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: local CPU embedding-model guardrail MVP (AISIX-Cloud#1331). +// +// The one acceptance path of the MVP vertical slice: a real request whose +// user text carries an EDA-software version number in natural-language +// Chinese goes through `/v1/chat/completions`, the in-process ONNX model +// judges the candidate's context window against the category prototype, +// and the version number is rewritten to `***`: +// - request side: the upstream's received body carries the masked text — +// the version number never left the gateway; +// - response side: the (fixed) upstream reply carrying the same sentence +// reaches the caller masked. +// +// SCOPE PINS (deliberate, per the MVP brief — not accidental gaps): +// - one happy path only; no negative/threshold/degrade cases; +// - non-streaming only: streamed output rides the guardrail's default +// BufferFull hold-back + the same segment pass, but is not pinned here; +// - /v1/chat/completions only: the sibling families (/v1/messages, +// /v1/responses, legacy completions, MCP) are explicitly unwired — +// tracked on the design issue, not silently missing. +// +// OPT-IN SPEC: skipped unless AISIX_LOCAL_GUARDRAIL_MODEL_DIR points at +// the model directory (model.onnx + tokenizer.json). Setting it implies +// the binary under test was built with `--features local-model-guardrail` +// (a default build would warn, serve unmasked, and fail this spec). +// The opt-in var deliberately carries the harness-stripped AISIX_ prefix +// so it can never leak into OTHER specs' spawned binaries; this spec +// forwards it explicitly as the binary's own GUARDRAIL_LOCAL_MODEL_DIR +// (non-AISIX on purpose — the config loader maps every AISIX_* env var +// onto a config field and strictly rejects unknown ones). +// Model files: https://huggingface.co/ibm-granite/granite-embedding-97m-multilingual-r2 +// (onnx/model_quint8_avx2.onnx saved as model.onnx, plus tokenizer.json). + +const CALLER = "sk-local-model-e2e-caller"; +const hash = (s: string) => createHash("sha256").update(s).digest("hex"); + +const SENSITIVE = "这个 EDA 软件的版本是 12.1"; +const MASKED = "这个 EDA 软件的版本是 ***"; + +const MODEL_DIR = process.env.AISIX_LOCAL_GUARDRAIL_MODEL_DIR; + +describe("local-model guardrail e2e: EDA version number masked on request and response", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let etcd: EtcdClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + if (!MODEL_DIR) return; + etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + // The mock reply carries the SAME sensitive sentence, so one request + // exercises both moderation hooks: input (what the upstream received) + // and output (what the caller got back). + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-local-model", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: SENSITIVE }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + + app = await spawnApp({ + // 2 lanes so the acceptance path exercises the session POOL + // dispatch (api7/aisix#1001), not just the single-lane degenerate + // case; behavior must be identical (lanes are stateless). + extraEnv: { + GUARDRAIL_LOCAL_MODEL_DIR: MODEL_DIR, + GUARDRAIL_LOCAL_MODEL_LANES: "2", + }, + }); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: "local-model-e2e-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "local-model-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + // Caller key last: it authenticating implies the whole seed set is in + // the DP snapshot (per this suite's readiness-gate rule). + await seed.createApiKey({ + key_hash: hash(CALLER), + allowed_models: ["local-model-e2e"], + }); + await waitConfigPropagation(async () => { + const r = await new ProxyClient(app!.proxyUrl, CALLER).listModels(); + return r.status === 200; + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test("version number becomes *** in the reply; the upstream never saw it", async (ctx) => { + if (!MODEL_DIR || !etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + const res = await new OpenAI({ + apiKey: CALLER, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }).chat.completions.create({ + model: "local-model-e2e", + messages: [{ role: "user", content: SENSITIVE }], + }); + + // Response side: the reply reaches the caller with the version number + // rewritten in place and everything else byte-identical. + expect(res.choices[0]?.message?.content).toBe(MASKED); + + // Request side: the upstream received the masked prompt — the version + // number never left the gateway. + const lastReq = upstream.receivedRequests.at(-1); + expect(lastReq).toBeDefined(); + expect(lastReq!.body).toContain(MASKED); + expect(lastReq!.body).not.toContain("12.1"); + }); +});