feat(guardrails): local CPU embedding-model guardrail MVP (AISIX-Cloud#1331)

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"

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"] }

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};

crates/aisix-proxy/src/chat.rs

@@ -1248,6 +1248,20 @@ async fn dispatch(
     *applied_out = applied_guardrails.clone();
     let resolved_chain: std::sync::Arc<dyn aisix_guardrails::Guardrail> =
         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<dyn aisix_guardrails::Guardrail> =
+        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

crates/aisix-proxy/src/state.rs

@@ -190,6 +190,13 @@ pub struct ProxyStateInner {
     pub health: Arc<HealthTracker>,
     /// Public liveness state served on `GET /livez`.
     pub livez: Arc<LivezState>,
+    /// 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<Arc<dyn aisix_guardrails::Guardrail>>,
     /// 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<dyn aisix_guardrails::Guardrail>,
+    ) -> 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.

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"

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)

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");
+  });
+});