@intflows/genkit-guard 0.0.7 → 0.0.8-alpha.1

package/README.md

@@ -1,4 +1,7 @@
-# **@intflows/genkit-guard**  
+# **@intflows/genkit-guard** 
+
+#### _This version uses OpenAI/privacy-filter instead of bert-base-NER_
+
 ### **Lightweight Intent, PII, and Safety Guardrails for Genkit**
 
 `@intflows/genkit-guard` provides a modular guardrail layer for Genkit flows.  
@@ -24,7 +27,7 @@ This library is designed for developers who want **practical, production‑ready
   Blocks jailbreak attempts using pattern‑based heuristics.
 
 - **Model‑Light Architecture**  
-  The package uses local `all-MiniLM-L6-v2` and `bert-base-NER` Models, these Models are downloaded once and cached locally.
+  The package uses local `all-MiniLM-L6-v2` and `openai/privacy-filter` Models, these Models are downloaded once and cached locally.
 
 - **Drop‑in Genkit Middleware**  
   Works with `ai.generate`, `ai.generateStream`, and Genkit flows.
@@ -38,12 +41,12 @@ This library is designed for developers who want **practical, production‑ready
 npm install @intflows/genkit-guard
 ```
 
-This library uses lightweight transformer models (MiniLM + BERT‑NER).  
+This library uses lightweight transformer models (MiniLM + Openai/privacy-filter).  
 
 Download them once. 
 
 ```bash
-## Download the transformer models (MiniLM + BERT‑NER)
+## Download the transformer models (MiniLM + OpenAI/privacy-filter)
 node node_modules/@intflows/genkit-guard/scripts/download-model.js
 ```
 
@@ -62,6 +65,7 @@ npm install @intflows/genkit-guard
 # Download Local Models (Only needed once)
 node node_modules/@intflows/genkit-guard/scripts/download-model.js
 ```
+_ This downloads the models to ./models folder, the total size is ~1.5 GB ( 1GB for Openai/privacy-filter + .5 GB for MiniLM-L6-v2) 
 
 ### 2. Update genkit
 
@@ -133,9 +137,9 @@ Before the LLM sees the prompt:
 Detected PII includes:
 
 - Emails  
-- Phone numbers  
-- Names (NER)  
+- Phone numbers    
 - AU identifiers (Medicare, TFN, ABN, etc.)
+- PII detected by local Model (OpenAI/privacy-filter)
 
 ### **3. LLM Call**
 The masked prompt is sent to the model.

package/dist/index.d.ts

@@ -3,4 +3,4 @@ export * from './core/types.js';
 /**
  * Pre-load the model to avoid cold-start delay on first user request.
  */
-export declare function initGuard(): Promise<void>;
+export declare function initGuard(config?: any): Promise<void>;

package/dist/index.js

@@ -5,11 +5,18 @@ export * from './core/types.js';
 /**
  * Pre-load the model to avoid cold-start delay on first user request.
  */
-export async function initGuard() {
+export async function initGuard(config) {
     console.log('[Guard] Loading local models...');
-    await Promise.all([
-        ModelSingleton.getExtractor(),
-        ModelSingleton.getNER()
-    ]);
+    const extractorModel = config?.models?.extractor ?? 'Xenova/all-MiniLM-L6-v2';
+    const piiModel = config?.pii?.model;
+    const piiMode = config?.pii?.mode ?? 'ner';
+    const tasks = [ModelSingleton.getExtractor(extractorModel)];
+    if (piiMode === 'ner') {
+        tasks.push(ModelSingleton.getNER(piiModel ?? 'Xenova/bert-base-NER'));
+    }
+    else {
+        tasks.push(ModelSingleton.getPIIClassifier(piiModel ?? 'openai/privacy-filter'));
+    }
+    await Promise.all(tasks);
     console.log('[Guard] Models loaded');
 }

package/dist/middleware/middleware.js

@@ -29,8 +29,12 @@ export function guard(config) {
         // -------------------------
         // 2. PII DETECTION + MASKING
         // -------------------------
-        const piiMatches = await detectPII(input);
-        console.log(`[PII Guard] Detected PII: ${piiMatches.length} matches found`);
+        const piiResponse = await detectPII(input, {
+            model: config?.pii?.model,
+            mode: config?.pii?.mode
+        });
+        const piiMatches = piiResponse?.matches || [];
+        console.log(`[PII Guard] Detected PII: ${piiMatches.length} matches found` + (piiResponse.classifier ? ' (classifier output present)' : ''));
         const tokenizer = new PiiTokenizer(); // <-- SINGLE INSTANCE
         const piiResult = tokenizer.mask(input, piiMatches);
         console.log(`[PII Guard] Masked PII: ${piiResult.piiTypes.length} types found`);
@@ -42,7 +46,10 @@ export function guard(config) {
             score: intentResult.score,
             piiDetected: piiMatches.length > 0,
             piiTypes: piiResult.piiTypes,
-            maskedInput: piiResult.maskedText
+            maskedInput: piiResult.maskedText,
+            piiModel: config?.pii?.model,
+            piiMode: config?.pii?.mode,
+            piiClassifierOutput: piiResponse.classifier
         };
         // Replace input
         req.prompt = piiResult.maskedText;

package/dist/pii/detector.d.ts

@@ -1,4 +1,10 @@
-export declare function detectPII(text: string): Promise<{
-    type: string;
-    value: string;
-}[]>;
+export declare function detectPII(text: string, opts?: {
+    model?: string;
+    mode?: 'ner' | 'classifier';
+}): Promise<{
+    matches: {
+        type: string;
+        value: string;
+    }[];
+    classifier: any;
+}>;

package/dist/pii/detector.js

@@ -15,20 +15,33 @@ const REGEX_RULES = [
     // CREDIT CARD (keep your existing one if needed)
     { type: 'CREDIT_CARD', pattern: /\b(?:\d[ -]*?){13,16}\b/g }
 ];
-export async function detectPII(text) {
-    const ner = await ModelSingleton.getNER();
+export async function detectPII(text, opts) {
+    const mode = opts?.mode ?? 'ner';
+    const model = opts?.model;
     const results = [];
-    // ---- NER ----
-    const entities = await ner(text);
-    for (const e of entities) {
-        if (e.entity.includes('PER')) {
-            results.push({ type: 'NAME', value: e.word.replace('##', '') });
-        }
-    }
-    // ---- REGEX ----
+    // ---- REGEX (always run) ----
     for (const rule of REGEX_RULES) {
         const matches = text.match(rule.pattern) || [];
         matches.forEach(m => results.push({ type: rule.type, value: m }));
     }
-    return results;
+    // ---- NER ----
+    let classifierOutput = undefined;
+    if (mode === 'ner') {
+        const ner = await ModelSingleton.getNER(model);
+        const entities = await ner(text);
+        for (const e of entities) {
+            if (e.entity && e.entity.includes('PER')) {
+                results.push({ type: 'NAME', value: (e.word || '').replace(/##/g, '') });
+            }
+        }
+    }
+    else {
+        // classifier mode: we call the classifier and return its output alongside regex matches.
+        const cls = await ModelSingleton.getPIIClassifier(model);
+        classifierOutput = await cls(text);
+    }
+    return {
+        matches: results,
+        classifier: classifierOutput
+    };
 }

package/dist/src/core/types.d.ts

@@ -0,0 +1,18 @@
+export interface IntentGuardConfig {
+    allowedIntent: string;
+    intents: Record<string, string>;
+    threshold?: number;
+    fallbackMessage?: string;
+}
+export interface IntentResult {
+    allowed: boolean;
+    score: number;
+}
+export interface PiiRule {
+    name: string;
+    pattern: RegExp;
+}
+export interface PiiConfig {
+    rules?: PiiRule[];
+    maskCharacter?: string;
+}

package/dist/src/core/types.js

@@ -0,0 +1 @@
+export {};

package/dist/src/index.d.ts

@@ -0,0 +1,6 @@
+export { guard } from './middleware/middleware.js';
+export * from './core/types.js';
+/**
+ * Pre-load the model to avoid cold-start delay on first user request.
+ */
+export declare function initGuard(config?: any): Promise<void>;

package/dist/src/index.js

@@ -0,0 +1,22 @@
+import { ModelSingleton } from './util/singleton.js';
+// export { intentGuard, piiGuard  } from './middleware/middleware.js';
+export { guard } from './middleware/middleware.js';
+export * from './core/types.js';
+/**
+ * Pre-load the model to avoid cold-start delay on first user request.
+ */
+export async function initGuard(config) {
+    console.log('[Guard] Loading local models...');
+    const extractorModel = config?.models?.extractor ?? 'Xenova/all-MiniLM-L6-v2';
+    const piiModel = config?.pii?.model;
+    const piiMode = config?.pii?.mode ?? 'ner';
+    const tasks = [ModelSingleton.getExtractor(extractorModel)];
+    if (piiMode === 'ner') {
+        tasks.push(ModelSingleton.getNER(piiModel ?? 'Xenova/bert-base-NER'));
+    }
+    else {
+        tasks.push(ModelSingleton.getPIIClassifier(piiModel ?? 'openai/privacy-filter'));
+    }
+    await Promise.all(tasks);
+    console.log('[Guard] Models loaded');
+}

package/dist/src/intent/intentAnalyzer.d.ts

@@ -0,0 +1,6 @@
+export declare function detectInjection(userInput: string): Promise<boolean>;
+export declare function analyzeIntentStructured(input: string, intents: Record<string, string>, threshold: number): Promise<{
+    intent: string;
+    score: number;
+    allowed: boolean;
+}>;

package/dist/src/intent/intentAnalyzer.js

@@ -0,0 +1,80 @@
+import { cos_sim } from '@huggingface/transformers';
+import { ModelSingleton } from '../util/singleton.js';
+const INJECTION_PATTERNS = [
+    'ignore instructions',
+    'reveal secrets',
+    'bypass security',
+    'system prompt',
+    'internal policy',
+    'developer message',
+    'confidential information',
+    'do not follow previous instructions',
+    'disregard earlier directives',
+    'override safety measures',
+    'leak data',
+    'expose vulnerabilities',
+    'access restricted areas',
+    'execute arbitrary code',
+    'perform unauthorized actions',
+    'escalate privileges',
+    'disable safeguards',
+    'manipulate output',
+    'circumvent controls',
+    'exploit weaknesses',
+    'compromise integrity',
+    'subvert protections',
+    'inject malicious content',
+    'craft deceptive prompts',
+    'trick the model',
+    'confuse the system',
+    'You are now a hacker',
+    'You are a malicious actor',
+    'Find a vulnerability',
+    'Write a prompt that bypasses security',
+    'Write a prompt that tricks the model',
+    'Write a prompt that confuses the system',
+    "ignore previous",
+    "ignore earlier",
+    "disregard previous",
+    "disregard earlier",
+    "you are a hacker",
+    "you are now a hacker",
+    "you are malicious",
+    "export the api key",
+    "give me the api key",
+    "bypass",
+    "override",
+    "jailbreak",
+    "system override",
+    "forget instructions",
+    "forget previous",
+    "forget earlier",
+    "do anything now",
+    "dan mode",
+    "developer mode",
+    "break character",
+    "act as"
+];
+export async function detectInjection(userInput) {
+    return INJECTION_PATTERNS.some(p => userInput.toLowerCase().includes(p));
+}
+export async function analyzeIntentStructured(input, intents, threshold) {
+    const extractor = await ModelSingleton.getExtractor();
+    let bestIntent = '';
+    let bestScore = 0;
+    for (const [key, desc] of Object.entries(intents)) {
+        const output = await extractor([`intent: ${desc}`, `intent: ${input}`], { pooling: 'mean', normalize: true });
+        const vectors = output.tolist();
+        const score = cos_sim(vectors[0], vectors[1]);
+        const finalScore = typeof score === 'number' ? score : score.data[0];
+        if (finalScore > bestScore) {
+            bestScore = finalScore;
+            bestIntent = key;
+        }
+    }
+    return {
+        intent: bestIntent,
+        score: bestScore,
+        allowed: bestScore >= threshold
+    };
+}

package/dist/src/middleware/middleware.d.ts

@@ -0,0 +1 @@
+export declare function guard(config: any): (req: any, next: any) => Promise<any>;

package/dist/src/middleware/middleware.js

@@ -0,0 +1,120 @@
+import { analyzeIntentStructured, detectInjection } from '../intent/intentAnalyzer.js';
+import { detectPII } from '../pii/detector.js';
+import { PiiTokenizer } from '../pii/tokenizer.js';
+export function guard(config) {
+    return async (req, next) => {
+        const input = req.prompt ||
+            req.messages?.[req.messages.length - 1]?.content?.[0]?.text ||
+            "";
+        // -------------------------
+        // 1. INTENT ANALYSIS
+        // -------------------------
+        const isInjection = await detectInjection(input);
+        if (isInjection) {
+            console.warn(`[Intent Guard] Prompt injection pattern detected in input`);
+            return block("Prompt injection detected", {
+                reason: "pattern_match"
+            });
+        }
+        console.log(`[Intent Guard] Analyzing intent for user's input`);
+        const intentResult = await analyzeIntentStructured(input, config.intent.semantic.intents, config.intent.semantic.threshold);
+        console.log(`[Intent Guard] Detected intent: ${intentResult.intent} (score: ${intentResult.score.toFixed(2)})`);
+        if (!intentResult.allowed) {
+            console.warn(`[Intent Guard] Intent "${intentResult.intent}" not allowed ${intentResult.allowed}`);
+            return block("Intent not allowed", {
+                intent: intentResult.intent,
+                score: intentResult.score
+            });
+        }
+        // -------------------------
+        // 2. PII DETECTION + MASKING
+        // -------------------------
+        const piiResponse = await detectPII(input, {
+            model: config?.pii?.model,
+            mode: config?.pii?.mode
+        });
+        const piiMatches = piiResponse?.matches || [];
+        console.log(`[PII Guard] Detected PII: ${piiMatches.length} matches found` + (piiResponse.classifier ? ' (classifier output present)' : ''));
+        const tokenizer = new PiiTokenizer(); // <-- SINGLE INSTANCE
+        const piiResult = tokenizer.mask(input, piiMatches);
+        console.log(`[PII Guard] Masked PII: ${piiResult.piiTypes.length} types found`);
+        // Attach tokenizer so response can unmask
+        req.metadata = {
+            ...req.metadata,
+            piiTokenizer: tokenizer,
+            intent: intentResult.intent,
+            score: intentResult.score,
+            piiDetected: piiMatches.length > 0,
+            piiTypes: piiResult.piiTypes,
+            maskedInput: piiResult.maskedText,
+            piiModel: config?.pii?.model,
+            piiMode: config?.pii?.mode,
+            piiClassifierOutput: piiResponse.classifier
+        };
+        // Replace input
+        req.prompt = piiResult.maskedText;
+        req.messages = [
+            {
+                role: 'user',
+                content: [{ text: piiResult.maskedText }],
+            },
+        ];
+        // -------------------------
+        // 3. LLM CALL
+        // -------------------------
+        const res = await next(req);
+        // -------------------------
+        // 4. RESPONSE UNMASK
+        // -------------------------
+        console.log(`[PII Guard] Unmasking response if needed`);
+        if (tokenizer) {
+            /**
+             * RECURSIVE TRANSFORMER
+             * This will find every string in the Genkit response (no matter if it's in
+             * candidates, message, custom, or output) and unmask it.
+             */
+            const transform = (obj) => {
+                // 1. If it's a string, unmask it
+                if (typeof obj === 'string') {
+                    return tokenizer.unmask(obj);
+                }
+                // 2. If it's an array, transform each element
+                if (Array.isArray(obj)) {
+                    return obj.map(transform);
+                }
+                // 3. If it's an object, transform each value
+                if (obj !== null && typeof obj === 'object') {
+                    // Note: We iterate keys and mutate the object directly 
+                    // to ensure Genkit's internal references are updated.
+                    for (const key of Object.keys(obj)) {
+                        obj[key] = transform(obj[key]);
+                    }
+                    return obj;
+                }
+                // 4. Return as-is for numbers/booleans/null
+                return obj;
+            };
+            // ----------------------------------------------------------------------
+            // 5. Transform the entire response object in-place to unmask all strings
+            // ----------------------------------------------------------------------      
+            transform(res);
+            console.log("[PII Guard] Deep unmasking complete across all candidates and custom fields.");
+        }
+        // ---------------------------------------------------------
+        // 6. Return the modified response with unmasked content
+        // ---------------------------------------------------------
+        return res;
+    };
+}
+// Helper to create a blocked response
+function block(message, metadata) {
+    return {
+        finishReason: 'blocked',
+        output: {
+            type: "error",
+            status: "BLOCKED",
+            message
+        },
+        metadata
+    };
+}

package/dist/src/pii/detector.d.ts

@@ -0,0 +1,10 @@
+export declare function detectPII(text: string, opts?: {
+    model?: string;
+    mode?: 'ner' | 'classifier';
+}): Promise<{
+    matches: {
+        type: string;
+        value: string;
+    }[];
+    classifier: any;
+}>;

package/dist/src/pii/detector.js

@@ -0,0 +1,47 @@
+import { ModelSingleton } from '../util/singleton.js';
+const REGEX_RULES = [
+    // EMAIL (keep your existing one)
+    { type: 'EMAIL', pattern: /\b[\w\.-]+@[\w\.-]+\.\w{2,}\b/gi },
+    // AU MOBILE (04xx xxx xxx or +61 4xx xxx xxx)
+    { type: 'AU_MOBILE', pattern: /\b(?:\+?61|0)4\d{2}[-\s]?\d{3}[-\s]?\d{3}\b/g },
+    // AU LANDLINE (02, 03, 07, 08)
+    { type: 'AU_LANDLINE', pattern: /\b(?:\+?61[-\s]?)?(?:2|3|7|8)\d{1}[-\s]?\d{4}[-\s]?\d{4}\b/g },
+    // MEDICARE NUMBER (10 digits, often grouped 4-5-1)
+    { type: 'MEDICARE', pattern: /\b\d{4}[-\s]?\d{5}[-\s]?\d\b/g },
+    // TFN (9 digits)
+    { type: 'TFN', pattern: /\b\d{3}[-\s]?\d{3}[-\s]?\d{3}\b/g },
+    // ABN (11 digits)
+    { type: 'ABN', pattern: /\b\d{2}[-\s]?\d{3}[-\s]?\d{3}[-\s]?\d{3}\b/g },
+    // CREDIT CARD (keep your existing one if needed)
+    { type: 'CREDIT_CARD', pattern: /\b(?:\d[ -]*?){13,16}\b/g }
+];
+export async function detectPII(text, opts) {
+    const mode = opts?.mode ?? 'ner';
+    const model = opts?.model;
+    const results = [];
+    // ---- REGEX (always run) ----
+    for (const rule of REGEX_RULES) {
+        const matches = text.match(rule.pattern) || [];
+        matches.forEach(m => results.push({ type: rule.type, value: m }));
+    }
+    // ---- NER ----
+    let classifierOutput = undefined;
+    if (mode === 'ner') {
+        const ner = await ModelSingleton.getNER(model);
+        const entities = await ner(text);
+        for (const e of entities) {
+            if (e.entity && e.entity.includes('PER')) {
+                results.push({ type: 'NAME', value: (e.word || '').replace(/##/g, '') });
+            }
+        }
+    }
+    else {
+        // classifier mode: we call the classifier and return its output alongside regex matches.
+        const cls = await ModelSingleton.getPIIClassifier(model);
+        classifierOutput = await cls(text);
+    }
+    return {
+        matches: results,
+        classifier: classifierOutput
+    };
+}

package/dist/src/pii/tokenizer.d.ts

@@ -0,0 +1,19 @@
+export type PiiResult = {
+    maskedText: string;
+    pii: Record<string, string>;
+    piiTypes: string[];
+};
+export declare class PiiTokenizer {
+    private vault;
+    private counter;
+    private piiTypes;
+    private createToken;
+    mask(text: string, matches: {
+        type: string;
+        value: string;
+    }[]): PiiResult;
+    unmask(text: string): string;
+    getVault(): {
+        [k: string]: string;
+    };
+}

package/dist/src/pii/tokenizer.js

@@ -0,0 +1,32 @@
+export class PiiTokenizer {
+    vault = new Map();
+    counter = 0;
+    piiTypes = new Set();
+    createToken(type) {
+        return `[[${type}_${this.counter++}]]`;
+    }
+    mask(text, matches) {
+        let masked = text;
+        for (const match of matches) {
+            const token = this.createToken(match.type);
+            this.vault.set(token, match.value);
+            this.piiTypes.add(match.type.toLowerCase());
+            masked = masked.split(match.value).join(token);
+        }
+        return {
+            maskedText: masked,
+            pii: Object.fromEntries(this.vault),
+            piiTypes: Array.from(this.piiTypes)
+        };
+    }
+    unmask(text) {
+        let result = text;
+        this.vault.forEach((value, token) => {
+            result = result.split(token).join(value); // Global replacement
+        });
+        return result;
+    }
+    getVault() {
+        return Object.fromEntries(this.vault);
+    }
+}

package/dist/src/util/singleton.d.ts

@@ -0,0 +1,14 @@
+export declare class ModelSingleton {
+    private static extractors;
+    private static nerClassifiers;
+    private static textClassifiers;
+    static init(): void;
+    static getExtractor(modelName?: string): Promise<any>;
+    static getNER(modelName?: string): Promise<any>;
+    static getPIIClassifier(modelName?: string): Promise<any>;
+    static preload(models?: {
+        extractor?: string;
+        ner?: string;
+        pii?: string;
+    }): Promise<void>;
+}

package/dist/src/util/singleton.js

@@ -0,0 +1,60 @@
+import { pipeline, env } from '@huggingface/transformers';
+import path from 'path';
+import fs from 'fs';
+import { fileURLToPath } from 'url';
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+env.allowRemoteModels = false;
+env.localModelPath = path.join(__dirname, '../../models');
+export class ModelSingleton {
+    static extractors = new Map();
+    static nerClassifiers = new Map();
+    static textClassifiers = new Map();
+    static init() {
+        // Always resolve model path relative to the client app, not the library
+        const projectRoot = process.cwd();
+        const modelPath = path.join(projectRoot, "models");
+        // Ensure directory exists
+        if (!fs.existsSync(modelPath)) {
+            fs.mkdirSync(modelPath, { recursive: true });
+        }
+        env.cacheDir = modelPath;
+        env.localModelPath = modelPath;
+        // Allow remote download if missing
+        env.allowRemoteModels = true;
+        console.log("[Guard] Using model directory:", modelPath);
+    }
+    static async getExtractor(modelName = 'Xenova/all-MiniLM-L6-v2') {
+        if (!this.extractors.has(modelName)) {
+            this.init();
+            const inst = await pipeline('feature-extraction', modelName);
+            this.extractors.set(modelName, inst);
+        }
+        return this.extractors.get(modelName);
+    }
+    static async getNER(modelName = 'Xenova/bert-base-NER') {
+        if (!this.nerClassifiers.has(modelName)) {
+            this.init();
+            const inst = await pipeline('token-classification', modelName);
+            this.nerClassifiers.set(modelName, inst);
+        }
+        return this.nerClassifiers.get(modelName);
+    }
+    static async getPIIClassifier(modelName = 'openai/privacy-filter') {
+        if (!this.textClassifiers.has(modelName)) {
+            this.init();
+            const inst = await pipeline('text-classification', modelName);
+            this.textClassifiers.set(modelName, inst);
+        }
+        return this.textClassifiers.get(modelName);
+    }
+    static async preload(models) {
+        const tasks = [];
+        if (models?.extractor)
+            tasks.push(this.getExtractor(models.extractor));
+        if (models?.ner)
+            tasks.push(this.getNER(models.ner));
+        if (models?.pii)
+            tasks.push(this.getPIIClassifier(models.pii));
+        await Promise.all(tasks);
+    }
+}

package/dist/util/singleton.d.ts

@@ -1,7 +1,14 @@
 export declare class ModelSingleton {
-    private static extractor;
-    private static classifier;
+    private static extractors;
+    private static nerClassifiers;
+    private static textClassifiers;
     static init(): void;
-    static getExtractor(): Promise<any>;
-    static getNER(): Promise<any>;
+    static getExtractor(modelName?: string): Promise<any>;
+    static getNER(modelName?: string): Promise<any>;
+    static getPIIClassifier(modelName?: string): Promise<any>;
+    static preload(models?: {
+        extractor?: string;
+        ner?: string;
+        pii?: string;
+    }): Promise<void>;
 }

package/dist/util/singleton.js

@@ -6,8 +6,9 @@ const __dirname = path.dirname(fileURLToPath(import.meta.url));
 env.allowRemoteModels = false;
 env.localModelPath = path.join(__dirname, '../../models');
 export class ModelSingleton {
-    static extractor = null;
-    static classifier = null;
+    static extractors = new Map();
+    static nerClassifiers = new Map();
+    static textClassifiers = new Map();
     static init() {
         // Always resolve model path relative to the client app, not the library
         const projectRoot = process.cwd();
@@ -22,19 +23,38 @@ export class ModelSingleton {
         env.allowRemoteModels = true;
         console.log("[Guard] Using model directory:", modelPath);
     }
-    static async getExtractor() {
-        if (!this.extractor) {
+    static async getExtractor(modelName = 'Xenova/all-MiniLM-L6-v2') {
+        if (!this.extractors.has(modelName)) {
             this.init();
-            this.extractor = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
+            const inst = await pipeline('feature-extraction', modelName);
+            this.extractors.set(modelName, inst);
         }
-        return this.extractor;
+        return this.extractors.get(modelName);
     }
-    static async getNER() {
-        if (!this.classifier) {
+    static async getNER(modelName = 'Xenova/bert-base-NER') {
+        if (!this.nerClassifiers.has(modelName)) {
             this.init();
-            // Token classification model for identifying sensitive entities
-            this.classifier = await pipeline('token-classification', 'Xenova/bert-base-NER');
+            const inst = await pipeline('token-classification', modelName);
+            this.nerClassifiers.set(modelName, inst);
         }
-        return this.classifier;
+        return this.nerClassifiers.get(modelName);
+    }
+    static async getPIIClassifier(modelName = 'openai/privacy-filter') {
+        if (!this.textClassifiers.has(modelName)) {
+            this.init();
+            const inst = await pipeline('text-classification', modelName);
+            this.textClassifiers.set(modelName, inst);
+        }
+        return this.textClassifiers.get(modelName);
+    }
+    static async preload(models) {
+        const tasks = [];
+        if (models?.extractor)
+            tasks.push(this.getExtractor(models.extractor));
+        if (models?.ner)
+            tasks.push(this.getNER(models.ner));
+        if (models?.pii)
+            tasks.push(this.getPIIClassifier(models.pii));
+        await Promise.all(tasks);
     }
 }

package/package.json

@@ -4,8 +4,25 @@
   "maintainers": [
     "Hemant Kohli"
   ],
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/IntFlows/genkit-guard"
+  },
+  "keywords": [
+    "genkit",
+    "genkit-guard",
+    "intflows",
+    "guardrails",
+    "intent-analysis",
+    "AI-safety",
+    "piimasking",
+    "content-moderation",
+    "NLP",
+    "transformers",
+    "huggingface"
+  ],
   "license": "Apache-2.0",
-  "version": "0.0.7",
+  "version": "0.0.8-alpha.1",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -20,15 +37,12 @@
     "build": "tsc",
     "prepublishOnly": "npm run build"
   },
-  "peerDependencies": {
-    "genkit": ">=0.5.0"
-  },
   "dependencies": {
-    "@huggingface/transformers": "^3.0.0",
-    "zod": "^3.23.0"
+    "@huggingface/transformers": "^4.2.0",
+    "zod": "^4.4.3"
   },
   "devDependencies": {
-    "typescript": "^5.6.0",
-    "@types/node": "^20.0.0"
+    "@types/node": "^25.8.0",
+    "typescript": "^6.0.3"
   }
 }

package/scripts/download-model.js

@@ -13,9 +13,14 @@ async function download() {
   await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2', {
     device: 'cpu'
   });
-  console.log('Downloading BERT-NER to ./models...');
-  await pipeline('token-classification', 'Xenova/bert-base-NER', {
-    device: 'cpu'
+  //console.log('Downloading BERT-NER to ./models...');
+  // await pipeline('token-classification', 'Xenova/bert-base-NER', {
+  //   device: 'cpu'
+  // });
+  
+  console.log('Downloading openai/privacy-filter to ./models...');
+  await pipeline('token-classification', 'openai/privacy-filter', {
+    device: 'cpu', dtype: "q4"
   });
   console.log('Model downloaded successfully.');
 }