ai-shield-classifier-onnx 0.4.0 → 0.5.1

package/dist/classifier.d.ts

@@ -0,0 +1,85 @@
+import type { Scanner, ScannerResult, ScanContext } from "ai-shield-core";
+/**
+ * Minimal subset of `onnxruntime-node`'s `InferenceSession` we use.
+ * Declaring it locally means this package type-checks even when
+ * `onnxruntime-node` is not installed (which is the entire point —
+ * it's an optional peer dep).
+ */
+export interface OnnxInferenceRuntime {
+    run(feeds: Record<string, OnnxTensorLike>): Promise<Record<string, OnnxTensorLike>>;
+    readonly inputNames?: readonly string[];
+    readonly outputNames?: readonly string[];
+}
+/** Tensor descriptor compatible with `onnxruntime-common`'s Tensor. */
+export interface OnnxTensorLike {
+    readonly data: ArrayLike<number> | BigInt64Array;
+    readonly dims: readonly number[];
+    readonly type?: string;
+}
+/**
+ * Tokenizer abstraction. Same trick as the runtime — we don't bind to
+ * any specific HF tokenizer package so the user can wire up
+ * `@huggingface/transformers`, `tokenizers`, or a hand-written one.
+ */
+export interface Tokenizer {
+    encode(text: string): {
+        input_ids: number[];
+        attention_mask: number[];
+        token_type_ids?: number[];
+    };
+    /** Optional max sequence length the tokenizer was trained for. */
+    modelMaxLength?: number;
+}
+export interface OnnxClassifierConfig {
+    /** A pre-constructed inference session. */
+    session: OnnxInferenceRuntime;
+    /** Tokenizer matching the model. */
+    tokenizer: Tokenizer;
+    /**
+     * Probability threshold above which the input is flagged as
+     * injection. Default 0.85 (calibrated for protectai/deberta-v3-base-
+     * prompt-injection — adjust per model).
+     */
+    threshold?: number;
+    /**
+     * Name of the output node that carries logits/probabilities.
+     * Default: first key in the runtime's result map.
+     */
+    outputName?: string;
+    /**
+     * Index of the "injection" class in the model output.
+     * Default: 1 (binary classifier convention: 0 = SAFE, 1 = INJECTION).
+     */
+    injectionClassIndex?: number;
+    /**
+     * Maximum sequence length to feed. Default: 512.
+     * Larger inputs are truncated head-only (start kept) which matches
+     * the standard DeBERTa fine-tune recipe.
+     */
+    maxLength?: number;
+}
+export declare class OnnxInjectionScanner implements Scanner {
+    readonly name = "onnx-classifier";
+    private readonly cfg;
+    constructor(config: OnnxClassifierConfig);
+    scan(input: string, _context: ScanContext): Promise<ScannerResult>;
+    /** Direct probability access — useful for tests + custom flows. */
+    predict(input: string): Promise<number>;
+}
+/**
+ * Convenience loader. Imports `onnxruntime-node` *at runtime* so
+ * consumers who never call this function don't pay the install cost.
+ *
+ * Tokenizer loading is left to the caller because tokenizer
+ * implementations vary widely between models — we don't want to
+ * pin a specific HF tokenizer package.
+ */
+export declare function loadOnnxClassifier(opts: {
+    modelPath: string;
+    tokenizer: Tokenizer;
+    threshold?: number;
+    outputName?: string;
+    injectionClassIndex?: number;
+    maxLength?: number;
+}): Promise<OnnxInjectionScanner>;
+//# sourceMappingURL=classifier.d.ts.map

package/dist/classifier.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"classifier.d.ts","sourceRoot":"","sources":["../src/classifier.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACV,OAAO,EACP,aAAa,EACb,WAAW,EAEZ,MAAM,gBAAgB,CAAC;AAqBxB;;;;;GAKG;AACH,MAAM,WAAW,oBAAoB;IACnC,GAAG,CACD,KAAK,EAAE,MAAM,CAAC,MAAM,EAAE,cAAc,CAAC,GACpC,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,cAAc,CAAC,CAAC,CAAC;IAC3C,QAAQ,CAAC,UAAU,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;IACxC,QAAQ,CAAC,WAAW,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;CAC1C;AAED,uEAAuE;AACvE,MAAM,WAAW,cAAc;IAC7B,QAAQ,CAAC,IAAI,EAAE,SAAS,CAAC,MAAM,CAAC,GAAG,aAAa,CAAC;IACjD,QAAQ,CAAC,IAAI,EAAE,SAAS,MAAM,EAAE,CAAC;IACjC,QAAQ,CAAC,IAAI,CAAC,EAAE,MAAM,CAAC;CACxB;AAED;;;;GAIG;AACH,MAAM,WAAW,SAAS;IACxB,MAAM,CAAC,IAAI,EAAE,MAAM,GAAG;QACpB,SAAS,EAAE,MAAM,EAAE,CAAC;QACpB,cAAc,EAAE,MAAM,EAAE,CAAC;QACzB,cAAc,CAAC,EAAE,MAAM,EAAE,CAAC;KAC3B,CAAC;IACF,kEAAkE;IAClE,cAAc,CAAC,EAAE,MAAM,CAAC;CACzB;AAED,MAAM,WAAW,oBAAoB;IACnC,2CAA2C;IAC3C,OAAO,EAAE,oBAAoB,CAAC;IAC9B,oCAAoC;IACpC,SAAS,EAAE,SAAS,CAAC;IACrB;;;;OAIG;IACH,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB;;;OAGG;IACH,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB;;;OAGG;IACH,mBAAmB,CAAC,EAAE,MAAM,CAAC;IAC7B;;;;OAIG;IACH,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB;AAED,qBAAa,oBAAqB,YAAW,OAAO;IAClD,QAAQ,CAAC,IAAI,qBAAqB;IAClC,OAAO,CAAC,QAAQ,CAAC,GAAG,CAEQ;gBAEhB,MAAM,EAAE,oBAAoB;IAkBlC,IAAI,CAAC,KAAK,EAAE,MAAM,EAAE,QAAQ,EAAE,WAAW,GAAG,OAAO,CAAC,aAAa,CAAC;IAmExE,mEAAmE;IAC7D,OAAO,CAAC,KAAK,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;CAkD9C;AAED;;;;;;;GAOG;AACH,wBAAsB,kBAAkB,CAAC,IAAI,EAAE;IAC7C,SAAS,EAAE,MAAM,CAAC;IAClB,SAAS,EAAE,SAAS,CAAC;IACrB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,UAAU,CAAC,EAAE,MAAM,CAAC;IACpB,mBAAmB,CAAC,EAAE,MAAM,CAAC;IAC7B,SAAS,CAAC,EAAE,MAAM,CAAC;CACpB,GAAG,OAAO,CAAC,oBAAoB,CAAC,CA6BhC"}

package/dist/classifier.js

@@ -0,0 +1,201 @@
+export class OnnxInjectionScanner {
+    name = "onnx-classifier";
+    cfg;
+    constructor(config) {
+        if (!config.session) {
+            throw new TypeError("OnnxInjectionScanner: 'session' is required");
+        }
+        if (!config.tokenizer) {
+            throw new TypeError("OnnxInjectionScanner: 'tokenizer' is required");
+        }
+        this.cfg = {
+            session: config.session,
+            tokenizer: config.tokenizer,
+            threshold: config.threshold ?? 0.85,
+            outputName: config.outputName,
+            injectionClassIndex: config.injectionClassIndex ?? 1,
+            maxLength: config.maxLength ?? config.tokenizer.modelMaxLength ?? 512,
+        };
+    }
+    async scan(input, _context) {
+        const start = performance.now();
+        try {
+            const probability = await this.predict(input);
+            const violations = [];
+            let decision = "allow";
+            if (probability >= this.cfg.threshold) {
+                decision = "block";
+                violations.push({
+                    type: "prompt_injection",
+                    scanner: this.name,
+                    score: probability,
+                    threshold: this.cfg.threshold,
+                    message: "ML classifier flagged prompt-injection",
+                    detail: `p(injection)=${probability.toFixed(4)} threshold=${this.cfg.threshold}`,
+                });
+            }
+            else if (probability >= this.cfg.threshold * 0.6) {
+                decision = "warn";
+                violations.push({
+                    type: "prompt_injection",
+                    scanner: this.name,
+                    score: probability,
+                    threshold: this.cfg.threshold,
+                    message: "ML classifier flagged borderline content",
+                    detail: `p(injection)=${probability.toFixed(4)} threshold=${this.cfg.threshold}`,
+                });
+            }
+            return {
+                decision,
+                violations,
+                durationMs: performance.now() - start,
+            };
+        }
+        catch (err) {
+            // ML errors must not take down the entire chain — degrade gracefully
+            // to "allow" with a synthetic violation so the audit log shows
+            // something went wrong without blocking traffic.
+            //
+            // Critic H4 round 1 — the raw error message can contain file
+            // paths (model location), native library symbols, or deployment-
+            // internal strings. Strip absolute paths before they hit the audit
+            // log. In dev mode we keep more detail to aid debugging.
+            const rawMessage = err?.message ?? "unknown error";
+            const isDev = process.env.NODE_ENV === "development" ||
+                process.env.AI_SHIELD_DEBUG === "1";
+            const safeDetail = isDev
+                ? rawMessage
+                : sanitizeOnnxErrorMessage(rawMessage);
+            return {
+                decision: "allow",
+                violations: [
+                    {
+                        type: "content_policy",
+                        scanner: this.name,
+                        score: 0,
+                        threshold: this.cfg.threshold,
+                        message: "ML classifier failed — degraded to allow",
+                        detail: safeDetail,
+                    },
+                ],
+                durationMs: performance.now() - start,
+            };
+        }
+    }
+    /** Direct probability access — useful for tests + custom flows. */
+    async predict(input) {
+        const tokens = this.cfg.tokenizer.encode(input);
+        const trunc = truncate(tokens, this.cfg.maxLength);
+        const inputIds = BigInt64Array.from(trunc.input_ids.map((n) => BigInt(n)));
+        const attentionMask = BigInt64Array.from(trunc.attention_mask.map((n) => BigInt(n)));
+        const dims = [1, trunc.input_ids.length];
+        const feeds = {
+            input_ids: { data: inputIds, dims, type: "int64" },
+            attention_mask: { data: attentionMask, dims, type: "int64" },
+        };
+        if (trunc.token_type_ids) {
+            feeds.token_type_ids = {
+                data: BigInt64Array.from(trunc.token_type_ids.map((n) => BigInt(n))),
+                dims,
+                type: "int64",
+            };
+        }
+        const result = await this.cfg.session.run(feeds);
+        const outputName = this.cfg.outputName ??
+            this.cfg.session.outputNames?.[0] ??
+            Object.keys(result)[0];
+        if (!outputName) {
+            throw new Error("OnnxInjectionScanner: no output node available");
+        }
+        const tensor = result[outputName];
+        if (!tensor) {
+            throw new Error(`OnnxInjectionScanner: output '${outputName}' not in result`);
+        }
+        // Model emits logits of shape [1, num_classes]. Softmax + pick class.
+        const logits = Array.from(tensor.data).map((n) => Number(n));
+        const probs = softmax(logits);
+        const idx = this.cfg.injectionClassIndex;
+        if (idx < 0 || idx >= probs.length) {
+            throw new Error(`OnnxInjectionScanner: injectionClassIndex ${idx} out of range (len=${probs.length})`);
+        }
+        return probs[idx] ?? 0;
+    }
+}
+/**
+ * Convenience loader. Imports `onnxruntime-node` *at runtime* so
+ * consumers who never call this function don't pay the install cost.
+ *
+ * Tokenizer loading is left to the caller because tokenizer
+ * implementations vary widely between models — we don't want to
+ * pin a specific HF tokenizer package.
+ */
+export async function loadOnnxClassifier(opts) {
+    // Dynamic import keeps the dep optional — TypeScript can't see it
+    // at compile time, which is the whole point.
+    let ort;
+    try {
+        ort = await import("onnxruntime-node");
+    }
+    catch (err) {
+        throw new Error("ai-shield-classifier-onnx: 'onnxruntime-node' is required " +
+            "to call loadOnnxClassifier(). Install it as a peer dependency.\n" +
+            `Underlying error: ${err.message}`);
+    }
+    const ortModule = ort;
+    const create = ortModule.InferenceSession?.create;
+    if (typeof create !== "function") {
+        throw new Error("ai-shield-classifier-onnx: 'onnxruntime-node' did not expose InferenceSession.create");
+    }
+    const session = await create(opts.modelPath);
+    return new OnnxInjectionScanner({
+        session,
+        tokenizer: opts.tokenizer,
+        threshold: opts.threshold,
+        outputName: opts.outputName,
+        injectionClassIndex: opts.injectionClassIndex,
+        maxLength: opts.maxLength,
+    });
+}
+// --- helpers ---
+/**
+ * Strip absolute paths and other deployment-internal strings from an
+ * ONNX-runtime error message before it lands in the audit log. Keeps
+ * the short error class / cause hint that helps diagnose the failure.
+ */
+function sanitizeOnnxErrorMessage(message) {
+    if (typeof message !== "string" || message.length === 0) {
+        return "classifier_runtime_error";
+    }
+    // Truncate before sanitizing — bounded work on adversarial input.
+    const truncated = message.length > 500 ? message.slice(0, 500) : message;
+    return truncated
+        // POSIX absolute paths
+        .replace(/(?:^|[\s(])(\/[\w./@-]+)/g, " [path]")
+        // Windows drive paths
+        .replace(/[A-Za-z]:\\[\\\w./@-]+/g, "[path]")
+        // file:// URLs
+        .replace(/file:\/\/\S+/g, "[file-url]")
+        // Memory addresses (0x...)
+        .replace(/0x[0-9a-fA-F]{6,}/g, "[addr]")
+        .trim();
+}
+function softmax(logits) {
+    if (logits.length === 0)
+        return [];
+    const max = Math.max(...logits);
+    const exps = logits.map((l) => Math.exp(l - max));
+    const sum = exps.reduce((a, b) => a + b, 0);
+    if (sum === 0)
+        return logits.map(() => 0);
+    return exps.map((e) => e / sum);
+}
+function truncate(tokens, maxLength) {
+    if (tokens.input_ids.length <= maxLength)
+        return tokens;
+    return {
+        input_ids: tokens.input_ids.slice(0, maxLength),
+        attention_mask: tokens.attention_mask.slice(0, maxLength),
+        token_type_ids: tokens.token_type_ids?.slice(0, maxLength),
+    };
+}
+//# sourceMappingURL=classifier.js.map

package/dist/classifier.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"classifier.js","sourceRoot":"","sources":["../src/classifier.ts"],"names":[],"mappings":"AA2FA,MAAM,OAAO,oBAAoB;IACtB,IAAI,GAAG,iBAAiB,CAAC;IACjB,GAAG,CAEQ;IAE5B,YAAY,MAA4B;QACtC,IAAI,CAAC,MAAM,CAAC,OAAO,EAAE,CAAC;YACpB,MAAM,IAAI,SAAS,CAAC,6CAA6C,CAAC,CAAC;QACrE,CAAC;QACD,IAAI,CAAC,MAAM,CAAC,SAAS,EAAE,CAAC;YACtB,MAAM,IAAI,SAAS,CAAC,+CAA+C,CAAC,CAAC;QACvE,CAAC;QACD,IAAI,CAAC,GAAG,GAAG;YACT,OAAO,EAAE,MAAM,CAAC,OAAO;YACvB,SAAS,EAAE,MAAM,CAAC,SAAS;YAC3B,SAAS,EAAE,MAAM,CAAC,SAAS,IAAI,IAAI;YACnC,UAAU,EAAE,MAAM,CAAC,UAAU;YAC7B,mBAAmB,EAAE,MAAM,CAAC,mBAAmB,IAAI,CAAC;YACpD,SAAS,EACP,MAAM,CAAC,SAAS,IAAI,MAAM,CAAC,SAAS,CAAC,cAAc,IAAI,GAAG;SAC7D,CAAC;IACJ,CAAC;IAED,KAAK,CAAC,IAAI,CAAC,KAAa,EAAE,QAAqB;QAC7C,MAAM,KAAK,GAAG,WAAW,CAAC,GAAG,EAAE,CAAC;QAChC,IAAI,CAAC;YACH,MAAM,WAAW,GAAG,MAAM,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC;YAC9C,MAAM,UAAU,GAAgB,EAAE,CAAC;YACnC,IAAI,QAAQ,GAA8B,OAAO,CAAC;YAElD,IAAI,WAAW,IAAI,IAAI,CAAC,GAAG,CAAC,SAAS,EAAE,CAAC;gBACtC,QAAQ,GAAG,OAAO,CAAC;gBACnB,UAAU,CAAC,IAAI,CAAC;oBACd,IAAI,EAAE,kBAAkB;oBACxB,OAAO,EAAE,IAAI,CAAC,IAAI;oBAClB,KAAK,EAAE,WAAW;oBAClB,SAAS,EAAE,IAAI,CAAC,GAAG,CAAC,SAAS;oBAC7B,OAAO,EAAE,wCAAwC;oBACjD,MAAM,EAAE,gBAAgB,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,cAAc,IAAI,CAAC,GAAG,CAAC,SAAS,EAAE;iBACjF,CAAC,CAAC;YACL,CAAC;iBAAM,IAAI,WAAW,IAAI,IAAI,CAAC,GAAG,CAAC,SAAS,GAAG,GAAG,EAAE,CAAC;gBACnD,QAAQ,GAAG,MAAM,CAAC;gBAClB,UAAU,CAAC,IAAI,CAAC;oBACd,IAAI,EAAE,kBAAkB;oBACxB,OAAO,EAAE,IAAI,CAAC,IAAI;oBAClB,KAAK,EAAE,WAAW;oBAClB,SAAS,EAAE,IAAI,CAAC,GAAG,CAAC,SAAS;oBAC7B,OAAO,EAAE,0CAA0C;oBACnD,MAAM,EAAE,gBAAgB,WAAW,CAAC,OAAO,CAAC,CAAC,CAAC,cAAc,IAAI,CAAC,GAAG,CAAC,SAAS,EAAE;iBACjF,CAAC,CAAC;YACL,CAAC;YAED,OAAO;gBACL,QAAQ;gBACR,UAAU;gBACV,UAAU,EAAE,WAAW,CAAC,GAAG,EAAE,GAAG,KAAK;aACtC,CAAC;QACJ,CAAC;QAAC,OAAO,GAAG,EAAE,CAAC;YACb,qEAAqE;YACrE,+DAA+D;YAC/D,iDAAiD;YACjD,EAAE;YACF,6DAA6D;YAC7D,iEAAiE;YACjE,mEAAmE;YACnE,yDAAyD;YACzD,MAAM,UAAU,GAAI,GAAa,EAAE,OAAO,IAAI,eAAe,CAAC;YAC9D,MAAM,KAAK,GACT,OAAO,CAAC,GAAG,CAAC,QAAQ,KAAK,aAAa;gBACtC,OAAO,CAAC,GAAG,CAAC,eAAe,KAAK,GAAG,CAAC;YACtC,MAAM,UAAU,GAAG,KAAK;gBACtB,CAAC,CAAC,UAAU;gBACZ,CAAC,CAAC,wBAAwB,CAAC,UAAU,CAAC,CAAC;YACzC,OAAO;gBACL,QAAQ,EAAE,OAAO;gBACjB,UAAU,EAAE;oBACV;wBACE,IAAI,EAAE,gBAAgB;wBACtB,OAAO,EAAE,IAAI,CAAC,IAAI;wBAClB,KAAK,EAAE,CAAC;wBACR,SAAS,EAAE,IAAI,CAAC,GAAG,CAAC,SAAS;wBAC7B,OAAO,EAAE,0CAA0C;wBACnD,MAAM,EAAE,UAAU;qBACnB;iBACF;gBACD,UAAU,EAAE,WAAW,CAAC,GAAG,EAAE,GAAG,KAAK;aACtC,CAAC;QACJ,CAAC;IACH,CAAC;IAED,mEAAmE;IACnE,KAAK,CAAC,OAAO,CAAC,KAAa;QACzB,MAAM,MAAM,GAAG,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC;QAChD,MAAM,KAAK,GAAG,QAAQ,CAAC,MAAM,EAAE,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,CAAC;QAEnD,MAAM,QAAQ,GAAG,aAAa,CAAC,IAAI,CAAC,KAAK,CAAC,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAC3E,MAAM,aAAa,GAAG,aAAa,CAAC,IAAI,CACtC,KAAK,CAAC,cAAc,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAC3C,CAAC;QACF,MAAM,IAAI,GAAG,CAAC,CAAC,EAAE,KAAK,CAAC,SAAS,CAAC,MAAM,CAAC,CAAC;QAEzC,MAAM,KAAK,GAAmC;YAC5C,SAAS,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE;YAClD,cAAc,EAAE,EAAE,IAAI,EAAE,aAAa,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE;SAC7D,CAAC;QACF,IAAI,KAAK,CAAC,cAAc,EAAE,CAAC;YACzB,KAAK,CAAC,cAAc,GAAG;gBACrB,IAAI,EAAE,aAAa,CAAC,IAAI,CAAC,KAAK,CAAC,cAAc,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;gBACpE,IAAI;gBACJ,IAAI,EAAE,OAAO;aACd,CAAC;QACJ,CAAC;QAED,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC;QACjD,MAAM,UAAU,GACd,IAAI,CAAC,GAAG,CAAC,UAAU;YACnB,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC;YACjC,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QACzB,IAAI,CAAC,UAAU,EAAE,CAAC;YAChB,MAAM,IAAI,KAAK,CAAC,gDAAgD,CAAC,CAAC;QACpE,CAAC;QACD,MAAM,MAAM,GAAG,MAAM,CAAC,UAAU,CAAC,CAAC;QAClC,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,MAAM,IAAI,KAAK,CACb,iCAAiC,UAAU,iBAAiB,CAC7D,CAAC;QACJ,CAAC;QAED,sEAAsE;QACtE,MAAM,MAAM,GAAG,KAAK,CAAC,IAAI,CAAC,MAAM,CAAC,IAAyB,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CACpE,MAAM,CAAC,CAAC,CAAC,CACV,CAAC;QACF,MAAM,KAAK,GAAG,OAAO,CAAC,MAAM,CAAC,CAAC;QAC9B,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,mBAAmB,CAAC;QACzC,IAAI,GAAG,GAAG,CAAC,IAAI,GAAG,IAAI,KAAK,CAAC,MAAM,EAAE,CAAC;YACnC,MAAM,IAAI,KAAK,CACb,6CAA6C,GAAG,sBAAsB,KAAK,CAAC,MAAM,GAAG,CACtF,CAAC;QACJ,CAAC;QACD,OAAO,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC;IACzB,CAAC;CACF;AAED;;;;;;;GAOG;AACH,MAAM,CAAC,KAAK,UAAU,kBAAkB,CAAC,IAOxC;IACC,kEAAkE;IAClE,6CAA6C;IAC7C,IAAI,GAAY,CAAC;IACjB,IAAI,CAAC;QACH,GAAG,GAAG,MAAM,MAAM,CAAC,kBAA4B,CAAC,CAAC;IACnD,CAAC;IAAC,OAAO,GAAG,EAAE,CAAC;QACb,MAAM,IAAI,KAAK,CACb,4DAA4D;YAC1D,kEAAkE;YAClE,qBAAsB,GAAa,CAAC,OAAO,EAAE,CAChD,CAAC;IACJ,CAAC;IACD,MAAM,SAAS,GAAI,GAA2F,CAAC;IAC/G,MAAM,MAAM,GAAG,SAAS,CAAC,gBAAgB,EAAE,MAAM,CAAC;IAClD,IAAI,OAAO,MAAM,KAAK,UAAU,EAAE,CAAC;QACjC,MAAM,IAAI,KAAK,CACb,sFAAsF,CACvF,CAAC;IACJ,CAAC;IACD,MAAM,OAAO,GAAG,MAAM,MAAM,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IAC7C,OAAO,IAAI,oBAAoB,CAAC;QAC9B,OAAO;QACP,SAAS,EAAE,IAAI,CAAC,SAAS;QACzB,SAAS,EAAE,IAAI,CAAC,SAAS;QACzB,UAAU,EAAE,IAAI,CAAC,UAAU;QAC3B,mBAAmB,EAAE,IAAI,CAAC,mBAAmB;QAC7C,SAAS,EAAE,IAAI,CAAC,SAAS;KAC1B,CAAC,CAAC;AACL,CAAC;AAED,kBAAkB;AAElB;;;;GAIG;AACH,SAAS,wBAAwB,CAAC,OAAe;IAC/C,IAAI,OAAO,OAAO,KAAK,QAAQ,IAAI,OAAO,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QACxD,OAAO,0BAA0B,CAAC;IACpC,CAAC;IACD,kEAAkE;IAClE,MAAM,SAAS,GAAG,OAAO,CAAC,MAAM,GAAG,GAAG,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC;IACzE,OAAO,SAAS;QACd,uBAAuB;SACtB,OAAO,CAAC,2BAA2B,EAAE,SAAS,CAAC;QAChD,sBAAsB;SACrB,OAAO,CAAC,yBAAyB,EAAE,QAAQ,CAAC;QAC7C,eAAe;SACd,OAAO,CAAC,eAAe,EAAE,YAAY,CAAC;QACvC,2BAA2B;SAC1B,OAAO,CAAC,oBAAoB,EAAE,QAAQ,CAAC;SACvC,IAAI,EAAE,CAAC;AACZ,CAAC;AAED,SAAS,OAAO,CAAC,MAAgB;IAC/B,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO,EAAE,CAAC;IACnC,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,GAAG,MAAM,CAAC,CAAC;IAChC,MAAM,IAAI,GAAG,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC,CAAC;IAClD,MAAM,GAAG,GAAG,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAC;IAC5C,IAAI,GAAG,KAAK,CAAC;QAAE,OAAO,MAAM,CAAC,GAAG,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC,CAAC;IAC1C,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,GAAG,CAAC,CAAC;AAClC,CAAC;AAED,SAAS,QAAQ,CACf,MAAuC,EACvC,SAAiB;IAEjB,IAAI,MAAM,CAAC,SAAS,CAAC,MAAM,IAAI,SAAS;QAAE,OAAO,MAAM,CAAC;IACxD,OAAO;QACL,SAAS,EAAE,MAAM,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,EAAE,SAAS,CAAC;QAC/C,cAAc,EAAE,MAAM,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC,EAAE,SAAS,CAAC;QACzD,cAAc,EAAE,MAAM,CAAC,cAAc,EAAE,KAAK,CAAC,CAAC,EAAE,SAAS,CAAC;KAC3D,CAAC;AACJ,CAAC"}

package/dist/index.d.ts

@@ -0,0 +1,2 @@
+export { OnnxInjectionScanner, loadOnnxClassifier, type OnnxClassifierConfig, type OnnxInferenceRuntime, type Tokenizer, } from "./classifier.js";
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAuBA,OAAO,EACL,oBAAoB,EACpB,kBAAkB,EAClB,KAAK,oBAAoB,EACzB,KAAK,oBAAoB,EACzB,KAAK,SAAS,GACf,MAAM,iBAAiB,CAAC"}

package/{src/index.ts → dist/index.js}

@@ -20,11 +20,5 @@
 // + tokenizer JSON, the wrapper handles inference + thresholding +
 // integration with the Scanner interface.
 // ============================================================
-
-export {
-  OnnxInjectionScanner,
-  loadOnnxClassifier,
-  type OnnxClassifierConfig,
-  type OnnxInferenceRuntime,
-  type Tokenizer,
-} from "./classifier.js";
+export { OnnxInjectionScanner, loadOnnxClassifier, } from "./classifier.js";
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,+DAA+D;AAC/D,6DAA6D;AAC7D,EAAE;AACF,yDAAyD;AACzD,kEAAkE;AAClE,gEAAgE;AAChE,EAAE;AACF,0BAA0B;AAC1B,yEAAyE;AACzE,mEAAmE;AACnE,uEAAuE;AACvE,gEAAgE;AAChE,EAAE;AACF,sBAAsB;AACtB,8DAA8D;AAC9D,oDAAoD;AACpD,mDAAmD;AACnD,EAAE;AACF,uEAAuE;AACvE,mEAAmE;AACnE,0CAA0C;AAC1C,+DAA+D;AAE/D,OAAO,EACL,oBAAoB,EACpB,kBAAkB,GAInB,MAAM,iBAAiB,CAAC"}

package/package.json

@@ -1,12 +1,12 @@
 {
   "name": "ai-shield-classifier-onnx",
-  "version": "0.4.0",
+  "version": "0.5.1",
   "license": "MIT",
   "description": "Optional ONNX ML classifier for ai-shield — DeBERTa-style prompt-injection detection alongside heuristic patterns",
   "author": "StudioMeyer <hello@studiomeyer.io>",
   "repository": {
     "type": "git",
-    "url": "https://github.com/studiomeyer-io/ai-shield",
+    "url": "git+https://github.com/studiomeyer-io/ai-shield.git",
     "directory": "packages/classifier-onnx"
   },
   "homepage": "https://github.com/studiomeyer-io/ai-shield/tree/main/packages/classifier-onnx",
@@ -24,6 +24,9 @@
     "ai-shield"
   ],
   "type": "module",
+  "files": [
+    "dist"
+  ],
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -37,7 +40,7 @@
     "typecheck": "tsc -b"
   },
   "peerDependencies": {
-    "ai-shield-core": "0.4.0",
+    "ai-shield-core": "0.5.0",
     "onnxruntime-node": ">=1.17.0"
   },
   "peerDependenciesMeta": {
@@ -46,7 +49,7 @@
     }
   },
   "devDependencies": {
-    "ai-shield-core": "0.4.0",
+    "ai-shield-core": "0.5.0",
     "typescript": "^5.7.0"
   },
   "engines": {

package/src/classifier.ts

@@ -1,326 +0,0 @@
-import type {
-  Scanner,
-  ScannerResult,
-  ScanContext,
-  Violation,
-} from "ai-shield-core";
-
-// ============================================================
-// OnnxInjectionScanner — ML-backed prompt-injection detection
-//
-// Implements the `Scanner` interface from ai-shield-core. Designed to
-// be added to a `ScannerChain` after the heuristic scanner so that
-// known patterns short-circuit on cheap regex AND novel paraphrases
-// still get a second-pass ML check.
-//
-// The runtime is abstracted via `OnnxInferenceRuntime` so this file
-// has zero hard dependency on `onnxruntime-node` at type-check time.
-// At runtime you pass in either:
-//   1. An already-constructed `ort.InferenceSession`
-//      (`new OnnxInjectionScanner({ session, tokenizer })`)
-//   2. A path to a `.onnx` model file + a path to a `tokenizer.json`
-//      (`await loadOnnxClassifier({ modelPath, tokenizerPath })`)
-//
-// Both paths keep the dep injection clean and unit-testable.
-// ============================================================
-
-/**
- * Minimal subset of `onnxruntime-node`'s `InferenceSession` we use.
- * Declaring it locally means this package type-checks even when
- * `onnxruntime-node` is not installed (which is the entire point —
- * it's an optional peer dep).
- */
-export interface OnnxInferenceRuntime {
-  run(
-    feeds: Record<string, OnnxTensorLike>,
-  ): Promise<Record<string, OnnxTensorLike>>;
-  readonly inputNames?: readonly string[];
-  readonly outputNames?: readonly string[];
-}
-
-/** Tensor descriptor compatible with `onnxruntime-common`'s Tensor. */
-export interface OnnxTensorLike {
-  readonly data: ArrayLike<number> | BigInt64Array;
-  readonly dims: readonly number[];
-  readonly type?: string;
-}
-
-/**
- * Tokenizer abstraction. Same trick as the runtime — we don't bind to
- * any specific HF tokenizer package so the user can wire up
- * `@huggingface/transformers`, `tokenizers`, or a hand-written one.
- */
-export interface Tokenizer {
-  encode(text: string): {
-    input_ids: number[];
-    attention_mask: number[];
-    token_type_ids?: number[];
-  };
-  /** Optional max sequence length the tokenizer was trained for. */
-  modelMaxLength?: number;
-}
-
-export interface OnnxClassifierConfig {
-  /** A pre-constructed inference session. */
-  session: OnnxInferenceRuntime;
-  /** Tokenizer matching the model. */
-  tokenizer: Tokenizer;
-  /**
-   * Probability threshold above which the input is flagged as
-   * injection. Default 0.85 (calibrated for protectai/deberta-v3-base-
-   * prompt-injection — adjust per model).
-   */
-  threshold?: number;
-  /**
-   * Name of the output node that carries logits/probabilities.
-   * Default: first key in the runtime's result map.
-   */
-  outputName?: string;
-  /**
-   * Index of the "injection" class in the model output.
-   * Default: 1 (binary classifier convention: 0 = SAFE, 1 = INJECTION).
-   */
-  injectionClassIndex?: number;
-  /**
-   * Maximum sequence length to feed. Default: 512.
-   * Larger inputs are truncated head-only (start kept) which matches
-   * the standard DeBERTa fine-tune recipe.
-   */
-  maxLength?: number;
-}
-
-export class OnnxInjectionScanner implements Scanner {
-  readonly name = "onnx-classifier";
-  private readonly cfg: Required<
-    Omit<OnnxClassifierConfig, "outputName">
-  > & { outputName?: string };
-
-  constructor(config: OnnxClassifierConfig) {
-    if (!config.session) {
-      throw new TypeError("OnnxInjectionScanner: 'session' is required");
-    }
-    if (!config.tokenizer) {
-      throw new TypeError("OnnxInjectionScanner: 'tokenizer' is required");
-    }
-    this.cfg = {
-      session: config.session,
-      tokenizer: config.tokenizer,
-      threshold: config.threshold ?? 0.85,
-      outputName: config.outputName,
-      injectionClassIndex: config.injectionClassIndex ?? 1,
-      maxLength:
-        config.maxLength ?? config.tokenizer.modelMaxLength ?? 512,
-    };
-  }
-
-  async scan(input: string, _context: ScanContext): Promise<ScannerResult> {
-    const start = performance.now();
-    try {
-      const probability = await this.predict(input);
-      const violations: Violation[] = [];
-      let decision: ScannerResult["decision"] = "allow";
-
-      if (probability >= this.cfg.threshold) {
-        decision = "block";
-        violations.push({
-          type: "prompt_injection",
-          scanner: this.name,
-          score: probability,
-          threshold: this.cfg.threshold,
-          message: "ML classifier flagged prompt-injection",
-          detail: `p(injection)=${probability.toFixed(4)} threshold=${this.cfg.threshold}`,
-        });
-      } else if (probability >= this.cfg.threshold * 0.6) {
-        decision = "warn";
-        violations.push({
-          type: "prompt_injection",
-          scanner: this.name,
-          score: probability,
-          threshold: this.cfg.threshold,
-          message: "ML classifier flagged borderline content",
-          detail: `p(injection)=${probability.toFixed(4)} threshold=${this.cfg.threshold}`,
-        });
-      }
-
-      return {
-        decision,
-        violations,
-        durationMs: performance.now() - start,
-      };
-    } catch (err) {
-      // ML errors must not take down the entire chain — degrade gracefully
-      // to "allow" with a synthetic violation so the audit log shows
-      // something went wrong without blocking traffic.
-      //
-      // Critic H4 round 1 — the raw error message can contain file
-      // paths (model location), native library symbols, or deployment-
-      // internal strings. Strip absolute paths before they hit the audit
-      // log. In dev mode we keep more detail to aid debugging.
-      const rawMessage = (err as Error)?.message ?? "unknown error";
-      const isDev =
-        process.env.NODE_ENV === "development" ||
-        process.env.AI_SHIELD_DEBUG === "1";
-      const safeDetail = isDev
-        ? rawMessage
-        : sanitizeOnnxErrorMessage(rawMessage);
-      return {
-        decision: "allow",
-        violations: [
-          {
-            type: "content_policy",
-            scanner: this.name,
-            score: 0,
-            threshold: this.cfg.threshold,
-            message: "ML classifier failed — degraded to allow",
-            detail: safeDetail,
-          },
-        ],
-        durationMs: performance.now() - start,
-      };
-    }
-  }
-
-  /** Direct probability access — useful for tests + custom flows. */
-  async predict(input: string): Promise<number> {
-    const tokens = this.cfg.tokenizer.encode(input);
-    const trunc = truncate(tokens, this.cfg.maxLength);
-
-    const inputIds = BigInt64Array.from(trunc.input_ids.map((n) => BigInt(n)));
-    const attentionMask = BigInt64Array.from(
-      trunc.attention_mask.map((n) => BigInt(n)),
-    );
-    const dims = [1, trunc.input_ids.length];
-
-    const feeds: Record<string, OnnxTensorLike> = {
-      input_ids: { data: inputIds, dims, type: "int64" },
-      attention_mask: { data: attentionMask, dims, type: "int64" },
-    };
-    if (trunc.token_type_ids) {
-      feeds.token_type_ids = {
-        data: BigInt64Array.from(trunc.token_type_ids.map((n) => BigInt(n))),
-        dims,
-        type: "int64",
-      };
-    }
-
-    const result = await this.cfg.session.run(feeds);
-    const outputName =
-      this.cfg.outputName ??
-      this.cfg.session.outputNames?.[0] ??
-      Object.keys(result)[0];
-    if (!outputName) {
-      throw new Error("OnnxInjectionScanner: no output node available");
-    }
-    const tensor = result[outputName];
-    if (!tensor) {
-      throw new Error(
-        `OnnxInjectionScanner: output '${outputName}' not in result`,
-      );
-    }
-
-    // Model emits logits of shape [1, num_classes]. Softmax + pick class.
-    const logits = Array.from(tensor.data as ArrayLike<number>).map((n) =>
-      Number(n),
-    );
-    const probs = softmax(logits);
-    const idx = this.cfg.injectionClassIndex;
-    if (idx < 0 || idx >= probs.length) {
-      throw new Error(
-        `OnnxInjectionScanner: injectionClassIndex ${idx} out of range (len=${probs.length})`,
-      );
-    }
-    return probs[idx] ?? 0;
-  }
-}
-
-/**
- * Convenience loader. Imports `onnxruntime-node` *at runtime* so
- * consumers who never call this function don't pay the install cost.
- *
- * Tokenizer loading is left to the caller because tokenizer
- * implementations vary widely between models — we don't want to
- * pin a specific HF tokenizer package.
- */
-export async function loadOnnxClassifier(opts: {
-  modelPath: string;
-  tokenizer: Tokenizer;
-  threshold?: number;
-  outputName?: string;
-  injectionClassIndex?: number;
-  maxLength?: number;
-}): Promise<OnnxInjectionScanner> {
-  // Dynamic import keeps the dep optional — TypeScript can't see it
-  // at compile time, which is the whole point.
-  let ort: unknown;
-  try {
-    ort = await import("onnxruntime-node" as string);
-  } catch (err) {
-    throw new Error(
-      "ai-shield-classifier-onnx: 'onnxruntime-node' is required " +
-        "to call loadOnnxClassifier(). Install it as a peer dependency.\n" +
-        `Underlying error: ${(err as Error).message}`,
-    );
-  }
-  const ortModule = (ort as { InferenceSession?: { create?: (path: string) => Promise<OnnxInferenceRuntime> } });
-  const create = ortModule.InferenceSession?.create;
-  if (typeof create !== "function") {
-    throw new Error(
-      "ai-shield-classifier-onnx: 'onnxruntime-node' did not expose InferenceSession.create",
-    );
-  }
-  const session = await create(opts.modelPath);
-  return new OnnxInjectionScanner({
-    session,
-    tokenizer: opts.tokenizer,
-    threshold: opts.threshold,
-    outputName: opts.outputName,
-    injectionClassIndex: opts.injectionClassIndex,
-    maxLength: opts.maxLength,
-  });
-}
-
-// --- helpers ---
-
-/**
- * Strip absolute paths and other deployment-internal strings from an
- * ONNX-runtime error message before it lands in the audit log. Keeps
- * the short error class / cause hint that helps diagnose the failure.
- */
-function sanitizeOnnxErrorMessage(message: string): string {
-  if (typeof message !== "string" || message.length === 0) {
-    return "classifier_runtime_error";
-  }
-  // Truncate before sanitizing — bounded work on adversarial input.
-  const truncated = message.length > 500 ? message.slice(0, 500) : message;
-  return truncated
-    // POSIX absolute paths
-    .replace(/(?:^|[\s(])(\/[\w./@-]+)/g, " [path]")
-    // Windows drive paths
-    .replace(/[A-Za-z]:\\[\\\w./@-]+/g, "[path]")
-    // file:// URLs
-    .replace(/file:\/\/\S+/g, "[file-url]")
-    // Memory addresses (0x...)
-    .replace(/0x[0-9a-fA-F]{6,}/g, "[addr]")
-    .trim();
-}
-
-function softmax(logits: number[]): number[] {
-  if (logits.length === 0) return [];
-  const max = Math.max(...logits);
-  const exps = logits.map((l) => Math.exp(l - max));
-  const sum = exps.reduce((a, b) => a + b, 0);
-  if (sum === 0) return logits.map(() => 0);
-  return exps.map((e) => e / sum);
-}
-
-function truncate(
-  tokens: ReturnType<Tokenizer["encode"]>,
-  maxLength: number,
-): ReturnType<Tokenizer["encode"]> {
-  if (tokens.input_ids.length <= maxLength) return tokens;
-  return {
-    input_ids: tokens.input_ids.slice(0, maxLength),
-    attention_mask: tokens.attention_mask.slice(0, maxLength),
-    token_type_ids: tokens.token_type_ids?.slice(0, maxLength),
-  };
-}

package/tsconfig.json

@@ -1,12 +0,0 @@
-{
-  "extends": "../../tsconfig.json",
-  "compilerOptions": {
-    "outDir": "./dist",
-    "rootDir": "./src",
-    "composite": true
-  },
-  "include": ["src/**/*"],
-  "references": [
-    { "path": "../core" }
-  ]
-}