diff --git a/crates/aisix-guardrails/src/local_model.rs b/crates/aisix-guardrails/src/local_model.rs index c62a7007..88d26785 100644 --- a/crates/aisix-guardrails/src/local_model.rs +++ b/crates/aisix-guardrails/src/local_model.rs @@ -1,23 +1,25 @@ -//! Local CPU embedding-model guardrail — MVP vertical slice of the -//! second-tier guardrail (AISIX-Cloud#1331). +//! Local CPU embedding-model guardrail — the second-tier guardrail +//! (AISIX-Cloud#1331), grown from the MVP vertical slice. //! -//! 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: +//! Implements the design issue's three-layer pipeline for one hardcoded +//! category (no prototype library resource, no standard risk categories +//! yet). 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. +//! 2. rule scoring ([`rules`]): hotword proximity co-occurrence raises +//! a candidate's score, negative patterns lower it, and a double +//! threshold resolves decisive candidates right here — high scores +//! rewrite and low scores release WITHOUT a model call; only the +//! uncertain band continues; +//! 3. a context window around each remaining 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 the category's prototype vector set (encoded at load +//! time; see [`PrototypeStrategy`]); above-threshold candidates are +//! rewritten in place to [`MASK_REPLACEMENT`]. +//! +//! Everything not rewritten is returned byte-identical. //! //! The verdict is always `Allow` — this guardrail rewrites, never blocks //! (the design issue's "只改写不阻断" hard constraint). @@ -90,6 +92,8 @@ //! ONNX Runtime download) with `ORT_LIB_PATH` pointing at a pre-fetched //! library (see `ort-sys` `build/vars.rs`). +mod rules; + use std::collections::BTreeMap; use std::ops::Range; use std::path::PathBuf; @@ -106,43 +110,135 @@ use crate::{ DEFAULT_STREAM_OUTPUT_BUFFER_BYTES, }; +use rules::{RuleDecision, RuleScorer}; + /// 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`]). +/// Optional cosine-similarity gate override (default: the configured +/// strategy's calibrated `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"; +/// Optional layer-② hotword proximity window override, in chars each +/// side of a candidate (default [`rules::DEFAULT_PROXIMITY_CHARS`], +/// clamped to [`rules::MAX_PROXIMITY_CHARS`]; malformed or zero → +/// default). Zero is a misconfiguration, not a mode — the same rule as +/// [`LANES_ENV`]: a zero window finds no hotword, so layer ② silently +/// stops masking while looking configured. +pub const RULE_WINDOW_ENV: &str = "GUARDRAIL_LOCAL_MODEL_RULE_WINDOW"; +/// Optional layer-③ prototype strategy: `description`, `max`, or +/// `centroid` (default [`PrototypeStrategy::default`]; malformed → +/// default with a warning — silently landing on the wrong vector space +/// would invalidate the operator's calibrated threshold). +pub const PROTOTYPES_ENV: &str = "GUARDRAIL_LOCAL_MODEL_PROTOTYPES"; /// 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. +/// The ONE category this module ships: EDA-software version numbers. +/// Under [`PrototypeStrategy::Description`] 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 软件的版本号"; +/// Sample sentences for the sample-based prototype strategies — the v2 +/// "customer supplies example sentences" path from the design issue, +/// collapsed to a compile-time constant set (synthesized; real customer +/// corpus not yet available). Coverage is by SHAPE, not by string: the +/// upgrade/rollback phrasing and the tool-name+version phrasing that the +/// single description prototype measurably missed, in Chinese and +/// English. Tool names and numbers are deliberately DIFFERENT from the +/// probe corpus (Spectre/Xcelium here, Virtuoso in the probes) so the +/// calibration probes measure shape generalization, not string overlap. +const PROTOTYPE_SAMPLES: &[&str] = &[ + "布局布线工具升级到 21.15 之后跑得快多了", + "仿真器回退到 19.03 才恢复正常", + "这个后端工具的版本是 33.0", + "综合工具的版本号是 2020.09,不要外传", + "Spectre 23.1.0 在这个工艺角下会崩溃", + "签核工具装的是 22.4 这个版本", + "We upgraded the place-and-route tool to 21.15", + "The simulator crashed on release 6.2.1", + "Xcelium 23.09 fails on this testbench", + "The sign-off tool version is 2020.09", +]; + +/// How the category's prototype vector set is built at load time. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PrototypeStrategy { + /// One vector: the embedded category description (the MVP form). + Description, + /// One vector per [`PROTOTYPE_SAMPLES`] entry; a window scores by its + /// MAX cosine over the set (nearest sample decides). + SampleMax, + /// One vector: the L2-renormalized mean of the sample embeddings; a + /// window scores against the class centroid. + SampleCentroid, +} + +impl PrototypeStrategy { + fn parse(raw: Option<&str>) -> Self { + let Some(raw) = raw else { + return Self::default(); + }; + match raw.to_ascii_lowercase().as_str() { + "description" => Self::Description, + "max" => Self::SampleMax, + "centroid" => Self::SampleCentroid, + other => { + tracing::warn!( + value = other, + default = ?Self::default(), + "unrecognized {PROTOTYPES_ENV}; using the default strategy" + ); + Self::default() + } + } + } + + /// Cosine gate calibrated per strategy with this module's + /// `#[ignore]` probe matrix (cosine absolute scale shifts with the + /// prototype construction, so one shared default would be wrong for + /// two of the three). Measured bands (granite-97m int8): + /// - `Description` keeps the MVP calibration: acceptance positive + /// ~0.90, every probed negative ≤0.76, hard positives 0.75–0.79 + /// below the gate — the measured single-prototype recall gap + /// (negative hard margin in all five phrasings swept) that + /// layer ② now covers. + /// - `SampleMax`: hard positives ≥0.8316, negatives ≤0.7867 + /// (hard margin +0.0449); 0.82 sits precision-leaning in that + /// band — 0.033 above the negative ceiling. + /// - `SampleCentroid`: hard positives ≥0.8616, negatives ≤0.8370 + /// (hard margin +0.0246); 0.85 likewise. + /// + /// All three lean precision — a mask false-positive corrupts user + /// content; layer ② carries recall for anchored shapes. + fn default_threshold(self) -> f32 { + match self { + Self::Description => 0.80, + Self::SampleMax => 0.82, + Self::SampleCentroid => 0.85, + } + } +} + +impl Default for PrototypeStrategy { + /// `SampleMax`: the probe matrix (module tests) measures the widest + /// positive hard margin here (+0.0449 vs +0.0246 for the centroid — + /// averaging ten shape-diverse samples into one vector costs + /// nearest-shape resolution). + fn default() -> Self { + Self::SampleMax + } +} + /// 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+)+"; @@ -157,6 +253,19 @@ const WINDOW_CONTEXT_CHARS: usize = 50; /// doing less, never to blocking or stalling. const MAX_MODEL_CALLS_PER_PASS: usize = 8; +/// Hard cap on rule-scored candidates per SEGMENT. Rule scoring is +/// µs-cheap per candidate but re-scans a proximity window each time, so +/// a crafted body that is nothing but candidates (`1.1 1.1 …`) turns +/// the per-segment scoring loop into a linear CPU amplifier on the +/// async worker — measured ~91 ms of synchronous work per MiB at the +/// default window and ~6× that at the window clamp (audit finding on +/// this PR; the request-body limit defaults to unlimited). Candidates +/// past the cap are RELEASED unscored with a warning — the same +/// fail-open arm as every other cap here — and a segment with thousands +/// of dotted-number runs is a log dump, not prose a version leaks +/// through. +const MAX_RULE_SCORED_SPANS_PER_SEGMENT: usize = 4096; + /// Hardcoded rewrite for an above-threshold candidate span. const MASK_REPLACEMENT: &str = "***"; @@ -194,28 +303,42 @@ pub struct LocalModelConfig { pub threshold: f32, /// Inference lanes = ONNX sessions = max concurrent inferences. pub lanes: usize, + /// Layer-② hotword proximity window (chars each side of a span). + pub rule_window: usize, + /// Layer-③ prototype-set construction. + pub prototypes: PrototypeStrategy, } 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 + /// malformed or out-of-range threshold falls back to the strategy's + /// 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). + /// configured-looking guardrail that silently never masks. Lanes and + /// the rule window follow the same lenient rule (malformed → + /// default). pub fn from_env() -> Option { let model_dir = PathBuf::from(std::env::var_os(MODEL_DIR_ENV)?); + let prototypes = PrototypeStrategy::parse(std::env::var(PROTOTYPES_ENV).ok().as_deref()); 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); + .unwrap_or_else(|| prototypes.default_threshold()); let lanes = parse_lanes(std::env::var(LANES_ENV).ok().as_deref()); + let rule_window = std::env::var(RULE_WINDOW_ENV) + .ok() + .and_then(|s| s.parse::().ok()) + .filter(|&w| w >= 1) + .unwrap_or(rules::DEFAULT_PROXIMITY_CHARS) + .min(rules::MAX_PROXIMITY_CHARS); Some(Self { model_dir, threshold, lanes, + rule_window, + prototypes, }) } } @@ -377,6 +500,35 @@ fn cosine(a: &[f32], b: &[f32]) -> f32 { a.iter().zip(b).map(|(x, y)| x * y).sum() } +/// A window's score against the prototype set: max cosine over the set +/// (with one vector this IS plain cosine, so all three strategies score +/// through here). +fn prototype_score(prototypes: &[Vec], v: &[f32]) -> f32 { + prototypes + .iter() + .map(|p| cosine(p, v)) + .fold(f32::NEG_INFINITY, f32::max) +} + +/// L2-renormalized mean of a set of L2-normalized vectors — the +/// [`PrototypeStrategy::SampleCentroid`] construction. +fn centroid(vectors: &[Vec]) -> Vec { + let dim = vectors.first().map_or(0, Vec::len); + let mut mean = vec![0.0f32; dim]; + for v in vectors { + for (m, x) in mean.iter_mut().zip(v) { + *m += x; + } + } + let norm = mean.iter().map(|v| v * v).sum::().sqrt(); + if norm > 0.0 { + for m in &mut mean { + *m /= norm; + } + } + mean +} + /// 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 @@ -430,10 +582,12 @@ fn apply_masks(text: &str, spans: &[Range]) -> String { /// The runtime guardrail. Always-`Allow`; masks via the segment hooks. pub struct LocalModelGuardrail { embedder: Arc, - /// L2-normalized embedding of [`PROTOTYPE_DESCRIPTION_ZH`]. - prototype: Vec, + /// L2-normalized prototype vector set (see [`PrototypeStrategy`]). + prototypes: Vec>, threshold: f32, candidate_re: Regex, + /// Layer-② scorer (hotword proximity + negative patterns). + rules: RuleScorer, /// 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. @@ -441,25 +595,40 @@ pub struct LocalModelGuardrail { } 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`. + /// Load tokenizer + the session pool and encode the category's + /// prototype set. Blocking (N model loads + up to + /// `PROTOTYPE_SAMPLES`-many inferences) — 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)?; + let embed_samples = || { + PROTOTYPE_SAMPLES + .iter() + .map(|s| embedder.embed(s)) + .collect::, _>>() + }; + let prototypes = match config.prototypes { + PrototypeStrategy::Description => vec![embedder.embed(PROTOTYPE_DESCRIPTION_ZH)?], + PrototypeStrategy::SampleMax => embed_samples()?, + PrototypeStrategy::SampleCentroid => vec![centroid(&embed_samples()?)], + }; tracing::info!( model_dir = %config.model_dir.display(), threshold = config.threshold, lanes = config.lanes, + rule_window = config.rule_window, + strategy = ?config.prototypes, + prototypes = prototypes.len(), load_ms = started.elapsed().as_millis() as u64, "local-model guardrail loaded (category: EDA software version)" ); Ok(Self { embedder: Arc::new(embedder), - prototype, + prototypes, threshold: config.threshold, candidate_re: Regex::new(CANDIDATE_PATTERN).expect("candidate pattern must compile"), + rules: RuleScorer::new(config.rule_window), permits: Arc::new(tokio::sync::Semaphore::new(config.lanes)), }) } @@ -489,38 +658,74 @@ impl LocalModelGuardrail { } /// 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 "只改写不阻断"). + /// were masked; `budget` is the shared per-pass model-call cap — + /// layer-② decisions are budget-free (µs of regex work), so the cap + /// only meters candidates that reach layer ③, and an exhausted + /// budget skips THOSE while later rule-decided candidates still + /// resolve. Model failure on a candidate leaves that span untouched + /// (rewrite less, never block — the fail-open arm of "只改写不阻断" + /// that also answers "model unavailable": the pipeline degrades to + /// ①+②). 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); - } + let (mut rule_masked, mut rule_passed, mut model_judged) = (0u32, 0u32, 0u32); + let mut over_budget = false; + let spans = candidate_spans(&self.candidate_re, text); + if spans.len() > MAX_RULE_SCORED_SPANS_PER_SEGMENT { + tracing::warn!( + candidates = spans.len(), + cap = MAX_RULE_SCORED_SPANS_PER_SEGMENT, + "local-model guardrail: candidate cap reached; the tail is released unscored" + ); + } + for span in spans.into_iter().take(MAX_RULE_SCORED_SPANS_PER_SEGMENT) { + match self.rules.decide(text, &span) { + RuleDecision::Mask => { + rule_masked += 1; + hits.push(span); } - Err(err) => { - tracing::warn!(error = %err, "local-model guardrail inference failed; span left unmasked"); + RuleDecision::Pass => rule_passed += 1, + RuleDecision::Model => { + if *budget == 0 { + if !over_budget { + over_budget = true; + tracing::warn!( + cap = MAX_MODEL_CALLS_PER_PASS, + "local-model guardrail: model-call cap reached; uncertain candidates left unmasked" + ); + } + continue; + } + *budget -= 1; + model_judged += 1; + let window = text[window_bounds(text, &span, WINDOW_CONTEXT_CHARS)].to_owned(); + match self.embed_window(window).await { + Ok(vector) => { + let score = prototype_score(&self.prototypes, &vector); + 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 rule_masked + rule_passed + model_judged > 0 { + tracing::debug!( + rule_masked, + rule_passed, + model_judged, + "local-model segment candidates resolved" + ); + } if hits.is_empty() { (text.to_owned(), 0) } else { @@ -650,12 +855,54 @@ mod tests { // parse fallback path via the public constructor contract. let cfg = LocalModelConfig { model_dir: PathBuf::from("/nonexistent"), - threshold: DEFAULT_THRESHOLD, + threshold: PrototypeStrategy::default().default_threshold(), lanes: 1, + rule_window: rules::DEFAULT_PROXIMITY_CHARS, + prototypes: PrototypeStrategy::default(), }; assert!(LocalModelGuardrail::load(&cfg).is_err()); } + #[test] + fn prototype_strategy_parses_and_defaults() { + assert_eq!(PrototypeStrategy::parse(None), PrototypeStrategy::default()); + assert_eq!( + PrototypeStrategy::parse(Some("description")), + PrototypeStrategy::Description + ); + assert_eq!( + PrototypeStrategy::parse(Some("max")), + PrototypeStrategy::SampleMax + ); + assert_eq!( + PrototypeStrategy::parse(Some("centroid")), + PrototypeStrategy::SampleCentroid + ); + // Case-insensitive: an operator typing `Description` means it. + assert_eq!( + PrototypeStrategy::parse(Some("Description")), + PrototypeStrategy::Description + ); + assert_eq!( + PrototypeStrategy::parse(Some("MAX")), + PrototypeStrategy::SampleMax + ); + assert_eq!( + PrototypeStrategy::parse(Some("bogus")), + PrototypeStrategy::default() + ); + } + + #[test] + fn centroid_is_the_renormalized_mean() { + let c = centroid(&[vec![1.0, 0.0], vec![0.0, 1.0]]); + let inv_sqrt2 = 1.0 / 2.0_f32.sqrt(); + assert!((c[0] - inv_sqrt2).abs() < 1e-6 && (c[1] - inv_sqrt2).abs() < 1e-6); + // Max-over-set picks the nearest prototype. + let set = vec![vec![1.0, 0.0], vec![0.0, 1.0]]; + assert!((prototype_score(&set, &[0.0, 1.0]) - 1.0).abs() < 1e-6); + } + #[test] fn parse_lanes_defaults_and_clamps() { // Unset / malformed / zero → 1 (zero is a misconfiguration, not @@ -680,55 +927,145 @@ mod tests { 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. + /// The MVP probe matrix, re-run for the prototype-set experiment: + /// the same 7 probe windows (1 acceptance-style positive, 2 hard + /// positives, 4 hard negatives) scored against 5 single-description + /// prototype phrasings (the MVP sweep that measured NEGATIVE margin + /// in every column) plus the two sample-based strategies. Prints the + /// full matrix and each column's hard margin + /// (min over hard positives − max over negatives). + /// + /// The MVP's original 5-phrasing sweep was scratch work; these + /// phrasings reconstruct it (the shipped description first) and are + /// committed so the experiment stays repeatable. #[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 phrasings = [ + PROTOTYPE_DESCRIPTION_ZH, + "软件版本号", + "芯片设计软件的版本号", + "提到了 EDA 工具的具体版本号", + "EDA 软件的版本信息,比如某个工具的版本是 12.1", ]; - let negatives = [ - "Elapsed: 12.345s, Memory: 4.2 GB", - "服务器的 IP 地址是 10.2.255.1", - "圆周率约等于 3.14159", - "工艺节点是 0.13um,良率还行", + let windows = [ + ("ACC", "这个 EDA 软件的版本是 12.1,请确认兼容性"), + ("POS", "我们把仿真工具升级到 2022.4 之后速度快了很多"), + ("POS", "Virtuoso IC6.1.8 出现了崩溃"), + ("NEG", "Elapsed: 12.345s, Memory: 4.2 GB"), + ("NEG", "服务器的 IP 地址是 10.2.255.1"), + ("NEG", "圆周率约等于 3.14159"), + ("NEG", "工艺节点是 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, + + // Columns: each phrasing as a single-vector prototype set, then + // the sample set (max) and its centroid. + let mut columns: Vec<(String, Vec>)> = Vec::new(); + for p in phrasings { + let v = g.embed_window(p.to_owned()).await.unwrap(); + columns.push((format!("desc:{p}"), vec![v])); + } + let mut samples = Vec::new(); + for s in PROTOTYPE_SAMPLES { + samples.push(g.embed_window((*s).to_owned()).await.unwrap()); + } + columns.push(("samples-max".to_owned(), samples.clone())); + columns.push(("samples-centroid".to_owned(), vec![centroid(&samples)])); + + for (name, prototypes) in &columns { + let mut hard_pos_min = f32::INFINITY; + let mut neg_max = f32::NEG_INFINITY; + println!("── column: {name}"); + for (kind, text) in windows { + let v = g.embed_window(text.to_owned()).await.unwrap(); + let s = prototype_score(prototypes, &v); + println!(" {kind} {s:.4} {text}"); + match kind { + "POS" => hard_pos_min = hard_pos_min.min(s), + "NEG" => neg_max = neg_max.max(s), + _ => {} + } + } + println!( + " hard margin (min POS − max NEG): {:+.4}", + hard_pos_min - neg_max ); - println!("POS(gap) {s:.4} {text}"); + + // Pin the calibration contract for the sample strategies: + // the margin the experiment claims stays open, and the + // shipped default gate sits strictly inside it. The + // description columns stay unasserted — their negative + // margin is the documented MVP finding, not a contract. + let gate = match name.as_str() { + "samples-max" => Some(PrototypeStrategy::SampleMax.default_threshold()), + "samples-centroid" => Some(PrototypeStrategy::SampleCentroid.default_threshold()), + _ => None, + }; + if let Some(gate) = gate { + assert!( + neg_max < gate && gate <= hard_pos_min, + "{name}: default gate {gate} outside the measured band ({neg_max:.4}, {hard_pos_min:.4}]" + ); + } } + } + + /// A candidate flood past [`MAX_RULE_SCORED_SPANS_PER_SEGMENT`] + /// releases the tail unscored: a rule-maskable sentence hidden + /// beyond the cap stays untouched (fail-open — rewrite less, never + /// stall), and the flood burns at most the model-call budget. + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn candidate_flood_releases_the_tail() { + let Some(g) = load_from_env() else { return }; + let flood = "1.1 ".repeat(MAX_RULE_SCORED_SPANS_PER_SEGMENT); + let text = format!("{flood}这个 EDA 软件的版本是 12.1"); + let outcome = g.moderate_input_segments(&[text]).await; + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); assert!( - acc >= g.threshold, - "acceptance positive {acc:.4} under threshold {}", - g.threshold + outcome.masked.is_none(), + "the over-cap tail must be released unscored" ); - for text in negatives { - let s = cosine( - &g.embed_window(text.to_owned()).await.unwrap(), - &g.prototype, - ); - println!("NEG {s:.4} {text}"); + } + + /// The acceptance matrix end to end through the segment hook, on the + /// default configuration: both hard positives and the MVP acceptance + /// positive are rewritten; every hard negative comes back untouched + /// (`masked == None` ⇒ byte-identical passthrough upstream). + #[tokio::test] + #[ignore = "needs GUARDRAIL_LOCAL_MODEL_DIR with model.onnx + tokenizer.json"] + async fn acceptance_matrix_end_to_end() { + let Some(g) = load_from_env() else { return }; + let masked_cases = [ + ("这个 EDA 软件的版本是 12.1", "这个 EDA 软件的版本是 ***"), + ( + "我们把仿真工具升级到 2022.4 之后速度快了很多", + "我们把仿真工具升级到 *** 之后速度快了很多", + ), + ("Virtuoso IC6.1.8 出现了崩溃", "Virtuoso IC*** 出现了崩溃"), + ]; + for (input, want) in masked_cases { + let outcome = g.moderate_input_segments(&[input.to_owned()]).await; + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); + let masked = outcome + .masked + .unwrap_or_else(|| panic!("must mask: {input}")); + assert_eq!(masked[0], want); + } + let passthrough_cases = [ + "Elapsed: 12.345s, Memory: 4.2 GB", + "服务器的 IP 地址是 10.2.255.1", + "圆周率约等于 3.14159", + "工艺节点是 0.13um,良率还行", + ]; + for input in passthrough_cases { + let outcome = g.moderate_input_segments(&[input.to_owned()]).await; + assert_eq!(outcome.verdict, GuardrailVerdict::Allow); assert!( - s < g.threshold, - "negative {s:.4} would mask at threshold {}: {text}", - g.threshold + outcome.masked.is_none(), + "negative must pass untouched: {input} → {:?}", + outcome.masked ); } } diff --git a/crates/aisix-guardrails/src/local_model/rules.rs b/crates/aisix-guardrails/src/local_model/rules.rs new file mode 100644 index 00000000..10270bae --- /dev/null +++ b/crates/aisix-guardrails/src/local_model/rules.rs @@ -0,0 +1,454 @@ +//! Layer ② of the three-layer pipeline (AISIX-Cloud#1331): rule scoring +//! between the regex candidate layer (①) and the model judgement (③). +//! +//! The MVP shipped ① → ③ directly, which put the whole separation burden +//! on zero-shot cosine — measured unable to split the hard positives +//! (`升级到 2022.4`, `Virtuoso IC6.1.8`, 0.75–0.79) from the compile-log +//! hard negatives (≤0.77). This layer restores the design's division of +//! labor: candidates with decisive lexical evidence are resolved here in +//! microseconds, and ONLY the genuinely ambiguous band pays a model call. +//! +//! Per candidate span: +//! - hotword co-occurrence RAISES the score. Three hotword classes +//! (Chinese triggers, English triggers, EDA tool names), each counted +//! once. A hotword ADJACENT to the span (≤ [`ADJACENT_GAP_CHARS`] +//! chars away — `版本是 12.1`, `Virtuoso IC6.1.8`) is decisive +//! evidence: +2 per class. Hotwords merely within the proximity +//! window are weak evidence: +1 total, regardless of class count — +//! distant co-occurrence alone must never mask, only route to the +//! model. +//! - negative patterns LOWER the score: -4 per class. A measurement +//! unit right after the number, an IPv4-shaped span, or a source +//! location is stronger evidence of "not a version" than any nearby +//! hotword is of the opposite, so one negative class (-4) outweighs +//! one decisive positive class (+2) — `版本升级后耗时 12.5s` must +//! not mask the timing. +//! - the double threshold turns the score into a [`RuleDecision`]: +//! ≥ [`MASK_SCORE`] rewrites without consulting the model, +//! ≤ [`PASS_SCORE`] releases without consulting the model, and only +//! the band in between (no evidence, or weak/conflicting evidence) +//! goes to layer ③. +//! +//! This is the mainstream DLP shape — hotword-proximity confidence +//! adjustment over a strong-format base pattern (Google Cloud DLP hotword +//! rules, Microsoft Purview supporting elements, AWS Macie +//! `maximumMatchDistance`, Palo Alto proximity keywords) — not an +//! invention of this crate. The proximity window default of +//! [`DEFAULT_PROXIMITY_CHARS`] sits at the tight end of the range those +//! engines use (Macie defaults to 50 chars, Palo Alto to 200, Purview +//! recommends 300, Google caps at 1000): the driving corpus is dense +//! compile logs where a wide window manufactures accidental +//! co-occurrence, and 50 also matches the ±50-char window layer ③ judges, +//! so the two layers reason about the same context. Override with +//! [`RULE_WINDOW_ENV`][super::RULE_WINDOW_ENV] (clamped to +//! [`MAX_PROXIMITY_CHARS`], the Google DLP hard cap). +//! +//! Threat-model boundary, inherited from the design issue's "只能防无意 +//! 泄漏" line and inherent to score-subtraction DLP: a sender (or a +//! hostile upstream, on the output side) can defeat the rule layer by +//! FORMATTING — appending a unit-looking suffix (`升级到 2022.4s`), a +//! `:digit` tail, or fullwidth digits that never become candidates. The +//! layer scores accidental phrasing, not adversarial encoding. +//! +//! Everything here is pure text work — no model, no I/O — so the layer is +//! unit-testable standalone, which is also how the "rules alone" halves +//! of the acceptance matrix are measured. + +use std::ops::Range; + +use regex::Regex; + +use super::window_bounds; + +/// Default hotword proximity window (chars each side of the span). +pub(super) const DEFAULT_PROXIMITY_CHARS: usize = 50; + +/// Upper clamp for the proximity window override. +pub(super) const MAX_PROXIMITY_CHARS: usize = 1000; + +/// A hotword within this many chars of the span edge is ADJACENT — +/// decisive rather than weak evidence. 8 chars covers the connective +/// tissue of every acceptance shape (`版本是 12.1` gap 2, `Virtuoso +/// IC6.1.8` gap 3, `upgrade to 21.15` gap 1, `version: v2022.4` gap 3) +/// while staying too small for an unrelated number to drift inside. +const ADJACENT_GAP_CHARS: usize = 8; + +/// Score for each hotword class with an adjacent match. +const ADJACENT_CLASS_SCORE: i32 = 2; +/// Score when hotwords exist only at window distance (flat, not per +/// class — distant co-occurrence stays weak no matter how many classes). +const WINDOW_EVIDENCE_SCORE: i32 = 1; +/// Score for each negative-pattern class that hits. +const NEGATIVE_CLASS_SCORE: i32 = -4; + +/// Decision bands: `score >= MASK_SCORE` masks without the model — one +/// adjacent hotword class alone is enough. +const MASK_SCORE: i32 = 2; +/// `score <= PASS_SCORE` releases without the model — one negative class +/// alone is enough, even against an adjacent hotword (-4 + 2). +const PASS_SCORE: i32 = -1; + +/// Chinese trigger hotwords (substring match — no word boundaries in +/// CJK). `版本号` is a substring of no other entry but kept explicit so +/// the list reads as the configured vocabulary. +const ZH_TRIGGERS: &[&str] = &["版本号", "版本", "升级到", "回退到"]; + +/// English trigger hotwords: word-bounded so `conversion` never fires +/// `version`. `upgraded? to` covers the bare and inflected forms. +const EN_TRIGGER_PATTERN: &str = r"(?i)\b(?:version|release|build|upgraded?\s+to)\b"; + +/// EDA tool names — the strongest anchors (`Virtuoso IC6.1.8` needs no +/// model). Word-bounded, case-insensitive (`vcs` / `VCS`). +const TOOL_PATTERN: &str = r"(?i)\b(?:virtuoso|calibre|vcs|innovus|icc2|primetime)\b"; + +/// Negative: a measurement unit right after the span (`12.345s`, +/// `4.2 GB`, `0.13um`, `99.9%`, `0.5ns`). Beyond the design brief's +/// minimum list, this covers the full timing-unit family (`us/ns/ps/fs`, +/// the `Hz` family, spelled-out durations) because the driving corpus — +/// STA/timing logs — is ns/ps-dense, and a unit the list misses next to +/// a tool name would RULE-MASK a slack value (audit finding on this PR). +/// Anchored to the span end with optional whitespace; letter units must +/// not continue into a longer word (`12.1 subsystem` is NOT an `s` hit), +/// checked with an explicit ASCII-alnum guard rather than `\b` because +/// the regex crate's Unicode `\b` treats a following CJK char as a word +/// char. +const UNIT_SUFFIX_PATTERN: &str = r"^\s*(?:%|(?i:ms|us|ns|ps|fs|s|secs?|seconds?|mins?|minutes?|hours?|[kmgt]i?b|um|nm|[kmg]?hz)(?:[^0-9A-Za-z]|$))"; + +/// Negative: the span itself is IPv4-shaped (`10.2.255.1`). Shape only — +/// no octet range check, matching how DLP engines treat dotted quads. +/// Known recall tradeoff (recorded on the design issue): real EDA +/// sub-versions can run to 4 dotted groups (`IC6.1.8.500`), and this +/// class releases them even next to an adjacent tool name. Kept anyway: +/// weakening the class when positive evidence is adjacent would mask +/// ACTUAL addresses (`Virtuoso 主机 10.2.255.1`), and a mask +/// false-positive corrupts content while a release only defers recall to +/// the sample corpus. +const IPV4_PATTERN: &str = r"^\d{1,3}(?:\.\d{1,3}){3}$"; + +/// Negative: source-location context — the span directly follows a +/// `file.ext:` prefix or is directly followed by `:digit` (`top.v:12.1`, +/// `12.1:3`-style diagnostics). +const FILE_COLON_PREFIX_PATTERN: &str = r"[\w.-]+\.[A-Za-z0-9]+:$"; +const COLON_DIGIT_SUFFIX_PATTERN: &str = r"^:\d"; + +/// What layer ② decided for one candidate span. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum RuleDecision { + /// Decisive positive evidence: rewrite, no model call. + Mask, + /// Decisive negative evidence: release, no model call. + Pass, + /// Uncertain band: only these consult layer ③. + Model, +} + +/// The compiled layer-② scorer. Hotword vocabulary and weights are +/// compile-time constants (the MVP's env-or-hardcoded config posture); +/// only the proximity window is operator-tunable. +pub(super) struct RuleScorer { + proximity_chars: usize, + en_trigger: Regex, + tool: Regex, + unit_suffix: Regex, + ipv4: Regex, + file_colon_prefix: Regex, + colon_digit_suffix: Regex, +} + +impl RuleScorer { + pub(super) fn new(proximity_chars: usize) -> Self { + let compile = |p: &str| Regex::new(p).expect("rule pattern must compile"); + Self { + // The env parse already rejects zero (a zero window finds no + // hotword — layer ② would silently stop masking), but the + // config fields are `pub`; clamp defensively like the lane + // pool does. + proximity_chars: proximity_chars.clamp(1, MAX_PROXIMITY_CHARS), + en_trigger: compile(EN_TRIGGER_PATTERN), + tool: compile(TOOL_PATTERN), + unit_suffix: compile(UNIT_SUFFIX_PATTERN), + ipv4: compile(IPV4_PATTERN), + file_colon_prefix: compile(FILE_COLON_PREFIX_PATTERN), + colon_digit_suffix: compile(COLON_DIGIT_SUFFIX_PATTERN), + } + } + + /// Score one candidate span of `text`. + pub(super) fn score(&self, text: &str, span: &Range) -> i32 { + let window = window_bounds(text, span, self.proximity_chars); + let mut adjacent_classes = 0i32; + let mut window_only = false; + + // Hotword classes: matches are searched inside the proximity + // window; a match within ADJACENT_GAP_CHARS of the span edge + // upgrades its class to decisive. + let mut tally = |ranges: &mut dyn Iterator>| { + let mut any = false; + let mut adjacent = false; + for r in ranges { + any = true; + if gap_chars(text, span, &r) <= ADJACENT_GAP_CHARS { + adjacent = true; + break; + } + } + if adjacent { + adjacent_classes += 1; + } else if any { + window_only = true; + } + }; + let win = &text[window.clone()]; + tally(&mut ZH_TRIGGERS.iter().flat_map(|t| { + win.match_indices(t) + .map(|(i, m)| window.start + i..window.start + i + m.len()) + })); + // A clipped window edge can bisect a word and hand `\b` a false + // boundary at the slice rim (`…conversion` clipped to a slice + // ending in `version`), so word-bounded matches flush with a + // CLIPPED edge are discarded. Substring (zh) matching has no + // boundary semantics, so it needs no such guard. + let clipped_start = window.start > 0; + let clipped_end = window.end < text.len(); + let mut bounded = |re: &Regex| { + tally( + &mut re + .find_iter(win) + .filter(|m| !(clipped_start && m.start() == 0)) + .filter(|m| !(clipped_end && m.end() == win.len())) + .map(|m| window.start + m.start()..window.start + m.end()), + ); + }; + bounded(&self.en_trigger); + bounded(&self.tool); + + let mut score = adjacent_classes * ADJACENT_CLASS_SCORE; + if adjacent_classes == 0 && window_only { + score += WINDOW_EVIDENCE_SCORE; + } + + // Negative classes: span-local shape checks, independent of the + // proximity window. + if self.unit_suffix.is_match(&text[span.end..]) { + score += NEGATIVE_CLASS_SCORE; + } + if self.ipv4.is_match(&text[span.clone()]) { + score += NEGATIVE_CLASS_SCORE; + } + if self.file_colon_prefix.is_match(&text[..span.start]) + || self.colon_digit_suffix.is_match(&text[span.end..]) + { + score += NEGATIVE_CLASS_SCORE; + } + score + } + + /// Apply the double threshold to [`score`](Self::score). + pub(super) fn decide(&self, text: &str, span: &Range) -> RuleDecision { + let score = self.score(text, span); + if score >= MASK_SCORE { + RuleDecision::Mask + } else if score <= PASS_SCORE { + RuleDecision::Pass + } else { + RuleDecision::Model + } + } +} + +/// Chars between a hotword match and the candidate span (0 when they +/// touch or overlap). Both ranges are byte ranges into `text`. +fn gap_chars(text: &str, span: &Range, hotword: &Range) -> usize { + if hotword.end <= span.start { + text[hotword.end..span.start].chars().count() + } else if hotword.start >= span.end { + text[span.end..hotword.start].chars().count() + } else { + 0 + } +} + +#[cfg(test)] +mod tests { + use super::super::{candidate_spans, CANDIDATE_PATTERN}; + use super::*; + + fn decisions(text: &str) -> Vec<(String, RuleDecision)> { + let scorer = RuleScorer::new(DEFAULT_PROXIMITY_CHARS); + let re = Regex::new(CANDIDATE_PATTERN).unwrap(); + candidate_spans(&re, text) + .into_iter() + .map(|s| (text[s.clone()].to_owned(), scorer.decide(text, &s))) + .collect() + } + + fn only(text: &str) -> RuleDecision { + let d = decisions(text); + assert_eq!(d.len(), 1, "expected one candidate in {text:?}: {d:?}"); + d[0].1 + } + + // ── the acceptance matrix, rules alone (no model) ──────────────────── + + #[test] + fn acceptance_positive_masks_by_rules_alone() { + assert_eq!( + only("这个 EDA 软件的版本是 12.1,请确认兼容性"), + RuleDecision::Mask + ); + } + + #[test] + fn hard_positives_mask_by_rules_alone() { + assert_eq!( + only("我们把仿真工具升级到 2022.4 之后速度快了很多"), + RuleDecision::Mask + ); + assert_eq!(only("Virtuoso IC6.1.8 出现了崩溃"), RuleDecision::Mask); + } + + #[test] + fn hard_negatives_pass_by_rules_alone() { + assert_eq!( + decisions("Elapsed: 12.345s, Memory: 4.2 GB"), + vec![ + ("12.345".to_owned(), RuleDecision::Pass), + ("4.2".to_owned(), RuleDecision::Pass), + ] + ); + assert_eq!(only("服务器的 IP 地址是 10.2.255.1"), RuleDecision::Pass); + assert_eq!(only("工艺节点是 0.13um,良率还行"), RuleDecision::Pass); + } + + #[test] + fn bare_number_stays_in_the_model_band() { + // No evidence either way — exactly what layer ③ exists for. + assert_eq!(only("圆周率约等于 3.14159"), RuleDecision::Model); + } + + // ── evidence weighting ─────────────────────────────────────────────── + + #[test] + fn negative_unit_outweighs_adjacent_trigger() { + // 版本/升级 hotwords adjacent, but the timing unit is decisive. + assert_eq!(only("版本升级后耗时 12.5s,可以接受"), RuleDecision::Pass); + } + + #[test] + fn window_distance_hotword_is_weak_evidence_only() { + // 版本 is inside the 50-char window but far from the number: + // weak evidence routes to the model instead of masking. + assert_eq!( + only("新版本已经发布了。另外今天集群的负载均值是 3.5"), + RuleDecision::Model + ); + } + + #[test] + fn english_triggers_are_word_bounded() { + assert_eq!(only("the conversion rate is 3.5"), RuleDecision::Model); + assert_eq!(only("we upgraded to 21.15 yesterday"), RuleDecision::Mask); + } + + #[test] + fn tool_name_is_decisive_only_when_adjacent() { + // Adjacent tool name: the design issue's "极强锚点" case. + assert_eq!(only("PrimeTime 2022.03 跑不过时序"), RuleDecision::Mask); + // Tool at window distance next to a non-version number: weak + // evidence — the model decides, not the rules. + assert_eq!( + only("PrimeTime reported a slack of 12.5"), + RuleDecision::Model + ); + } + + #[test] + fn clipped_window_edge_cannot_fabricate_a_word_boundary() { + // With a window sized to clip `conversion` exactly at its inner + // `version`, the slice-start `\b` would fire and rule-mask; the + // clipped-edge guard drops the match and the candidate stays in + // the model band. + let text = "big conversion 3.5 result"; + let re = Regex::new(CANDIDATE_PATTERN).unwrap(); + let span = &candidate_spans(&re, text)[0]; + let tight = RuleScorer::new(8); + assert_eq!(tight.decide(text, span), RuleDecision::Model); + } + + #[test] + fn source_location_context_passes() { + assert_eq!( + only("see top.v:12.1 for the assignment"), + RuleDecision::Pass + ); + } + + #[test] + fn unit_must_not_continue_into_a_word() { + // `s` starts `subsystem`: not a unit hit, and 版本 is adjacent → + // decisive positive stands. + assert_eq!( + only("版本是 12.1 subsystem 之外的另一个话题"), + RuleDecision::Mask + ); + } + + #[test] + fn timing_units_outweigh_adjacent_tool_names() { + // The driving corpus is ns/ps-dense STA logs: a slack value next + // to a tool name must NOT rule-mask (audit finding on this PR). + assert_eq!(only("PrimeTime slack 0.5ns 违例"), RuleDecision::Pass); + assert_eq!(only("版本升级后耗时 12.5ns,可以接受"), RuleDecision::Pass); + assert_eq!(only("clock period 1.25ps setup ok"), RuleDecision::Pass); + assert_eq!(only("跑到 3.2GHz 依然稳定"), RuleDecision::Pass); + assert_eq!(only("内存占用 1.5 GiB 左右"), RuleDecision::Pass); + } + + #[test] + fn spelled_duration_units_are_negative_evidence() { + assert_eq!(only("the build took 12.5 minutes"), RuleDecision::Pass); + assert_eq!(only("版本是 12.1 seconds 之外的话题"), RuleDecision::Pass); + } + + #[test] + fn proximity_window_bounds_the_hotword_search() { + // Same sentence, tool name pushed outside a tiny window: the + // evidence disappears and the candidate falls to the model band. + let text = "Innovus 的运行日志我贴在下面了,请帮忙看看统计值 21.12"; + let tight = RuleScorer::new(4); + let wide = RuleScorer::new(DEFAULT_PROXIMITY_CHARS); + let re = Regex::new(CANDIDATE_PATTERN).unwrap(); + let span = &candidate_spans(&re, text)[0]; + assert_eq!(tight.decide(text, span), RuleDecision::Model); + assert_eq!(wide.decide(text, span), RuleDecision::Model); + // At window distance the tool is weak (+1) evidence, not a mask. + assert_eq!(wide.score(text, span), 1); + assert_eq!(tight.score(text, span), 0); + } + + // ── report instrument: model-call rate over the probe corpus ──────── + + /// The layer-③ call-rate accounting the design issue asks for: over + /// the MVP probe corpus (7 windows, 8 candidates), layer ② resolves + /// 7 of 8 candidates and only the bare-number window still pays a + /// model call — versus 8 of 8 in the ①→③ MVP. + #[test] + fn probe_corpus_model_call_rate() { + let corpus = [ + "这个 EDA 软件的版本是 12.1,请确认兼容性", + "我们把仿真工具升级到 2022.4 之后速度快了很多", + "Virtuoso IC6.1.8 出现了崩溃", + "Elapsed: 12.345s, Memory: 4.2 GB", + "服务器的 IP 地址是 10.2.255.1", + "圆周率约等于 3.14159", + "工艺节点是 0.13um,良率还行", + ]; + let all: Vec = corpus + .iter() + .flat_map(|t| decisions(t).into_iter().map(|(_, d)| d)) + .collect(); + assert_eq!(all.len(), 8); + let to_model = all.iter().filter(|d| **d == RuleDecision::Model).count(); + assert_eq!(to_model, 1, "decisions: {all:?}"); + } +} diff --git a/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts b/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts index 6ae0ef43..3e526b40 100644 --- a/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts +++ b/tests/e2e/src/cases/guardrail-local-model-e2e.test.ts @@ -12,20 +12,21 @@ import { type SpawnedApp, } from "../harness/index.js"; -// E2E: local CPU embedding-model guardrail MVP (AISIX-Cloud#1331). +// E2E: local CPU embedding-model guardrail (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. +// Two acceptance paths through `/v1/chat/completions`, each asserted on +// both sides (request: what the upstream received; response: what the +// caller got back): +// - the MVP acceptance sentence: an EDA-software version number in +// natural-language Chinese is rewritten to `***`; +// - the layer-② acceptance matrix: one message carrying both hard +// positives (`升级到 2022.4`, `Virtuoso IC6.1.8` — the shapes the MVP +// measurably missed) and the hard negatives (compile-log timing, +// memory size, IPv4, bare number); the positives are rewritten, the +// negatives come back byte-identical, and the bare number traverses a +// live layer-③ inference (no lexical evidence → model band). // // 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, @@ -50,39 +51,55 @@ const hash = (s: string) => createHash("sha256").update(s).digest("hex"); const SENSITIVE = "这个 EDA 软件的版本是 12.1"; const MASKED = "这个 EDA 软件的版本是 ***"; +// The layer-② acceptance matrix in one message: both hard positives, the +// hard negatives, and a bare number with no lexical evidence at all (the +// model band — this candidate pays a real in-process inference). +const MIXED = + "我们把仿真工具升级到 2022.4 之后,Virtuoso IC6.1.8 反而开始频繁崩溃," + + "完整的运行日志我贴在下面了,麻烦帮忙看看到底是哪一步出了问题: " + + "Elapsed: 12.345s, Memory: 4.2 GB, 服务器 IP 是 10.2.255.1, " + + "另外圆周率约等于 3.14159"; +const MIXED_MASKED = + "我们把仿真工具升级到 *** 之后,Virtuoso IC*** 反而开始频繁崩溃," + + "完整的运行日志我贴在下面了,麻烦帮忙看看到底是哪一步出了问题: " + + "Elapsed: 12.345s, Memory: 4.2 GB, 服务器 IP 是 10.2.255.1, " + + "另外圆周率约等于 3.14159"; + 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 mixedUpstream: OpenAiUpstream | undefined; let etcd: EtcdClient | undefined; let etcdReachable = false; + const reply = (content: string) => ({ + id: "cmpl-local-model", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }); + 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 }, - }, - }); + // Each mock reply carries the SAME text as its test's prompt, so one + // request exercises both moderation hooks: input (what the upstream + // received) and output (what the caller got back). + upstream = await startOpenAiUpstream({ nonStreamBody: reply(SENSITIVE) }); + mixedUpstream = await startOpenAiUpstream({ nonStreamBody: reply(MIXED) }); app = await spawnApp({ // 2 lanes so the acceptance path exercises the session POOL @@ -106,11 +123,22 @@ describe("local-model guardrail e2e: EDA version number masked on request and re model_name: "gpt-4o-mini", provider_key_id: pk.id, }); + const mixedPk = await seed.createProviderKey({ + display_name: "local-model-e2e-mixed-pk", + secret: "sk-mock", + api_base: `${mixedUpstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "local-model-e2e-mixed", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: mixedPk.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"], + allowed_models: ["local-model-e2e", "local-model-e2e-mixed"], }); await waitConfigPropagation(async () => { const r = await new ProxyClient(app!.proxyUrl, CALLER).listModels(); @@ -121,6 +149,7 @@ describe("local-model guardrail e2e: EDA version number masked on request and re afterAll(async () => { await app?.exit(); await upstream?.close(); + await mixedUpstream?.close(); }); test("version number becomes *** in the reply; the upstream never saw it", async (ctx) => { @@ -149,4 +178,32 @@ describe("local-model guardrail e2e: EDA version number masked on request and re expect(lastReq!.body).toContain(MASKED); expect(lastReq!.body).not.toContain("12.1"); }); + + test("hard positives are rewritten while hard negatives pass byte-identical", async (ctx) => { + if (!MODEL_DIR || !etcdReachable || !app || !mixedUpstream) { + ctx.skip(); + return; + } + + const res = await new OpenAI({ + apiKey: CALLER, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }).chat.completions.create({ + model: "local-model-e2e-mixed", + messages: [{ role: "user", content: MIXED }], + }); + + // Response side: versions masked, everything else — the compile-log + // numbers, the IP, the bare number the model judged — byte-identical. + expect(res.choices[0]?.message?.content).toBe(MIXED_MASKED); + + // Request side: the upstream saw the same rewrite and neither + // version value. + const lastReq = mixedUpstream.receivedRequests.at(-1); + expect(lastReq).toBeDefined(); + expect(lastReq!.body).toContain(MIXED_MASKED); + expect(lastReq!.body).not.toContain("2022.4"); + expect(lastReq!.body).not.toContain("6.1.8"); + }); });