@intflows/genkit-guard 0.0.4

package/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

package/README.md

@@ -0,0 +1,203 @@
+# **@intflows/genkit-guard**  
+### **Lightweight Intent, PII, and Safety Guardrails for Genkit**
+
+`@intflows/genkit-guard` provides a modular guardrail layer for Genkit flows.  
+It adds **semantic intent validation**, **PII masking/unmasking**, and **prompt‑injection detection** with minimal configuration.
+
+This library is designed for developers who want **practical, production‑ready safety controls** without heavy dependencies or complex setup.
+
+---
+
+## ✨ Features
+
+- **Semantic Intent Guarding**  
+  Uses MiniLM embeddings to ensure prompts match allowed intents.
+
+- **PII Detection & Masking**  
+  Detects emails, phone numbers, names, and AU‑specific identifiers.  
+  Replaces PII with reversible tokens before sending to the LLM.
+
+- **Automatic Unmasking**  
+  Restores original PII in the model’s response, even inside structured JSON.
+
+- **Prompt Injection Detection**  
+  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.
+
+- **Drop‑in Genkit Middleware**  
+  Works with `ai.generate`, `ai.generateStream`, and Genkit flows.
+
+---
+
+### 📦 Installation
+
+```bash
+## Install the package
+npm install @intflows/genkit-guard
+```
+
+This library uses lightweight transformer models (MiniLM + BERT‑NER).  
+
+Download them once. 
+
+```bash
+## Download the transformer models (MiniLM + BERT‑NER)
+node node_modules/@intflows/genkit-guard/scripts/download-model.js
+```
+
+Models are cached locally and reused across runs.
+
+---
+
+## 🚀 Quick Start
+
+### 1. Initialize Local folder
+
+```bash
+# Install @intflows/genkit-guard
+npm install @intflows/genkit-guard
+
+# Download Local Models (Only needed once)
+node node_modules/@intflows/genkit-guard/scripts/download-model.js
+```
+
+### 2. Update genkit
+
+```ts
+import { guard, initGuard } from "@intflows/genkit-guard";
+
+await initGuard();
+
+const response = await ai.generate({
+  prompt: "How do I integrate with Azure Blob Storage?",
+  use: [
+    guard({
+      intent: {
+        mode: "semantic",
+        allowedIntent: "integration",
+        semantic: {
+          threshold: 0.7,
+          intents: {
+            integration: "Azure Blob, APIs, workflows"
+          }
+        }
+      },
+      pii: { reversible: true }
+    })
+  ]
+});
+```
+
+**You Can also check the full step by step guide here:**
+
+[Intflows Wiki](https://github.com/IntFlows/genkit-guard/wiki)
+
+### 3. Execute the Genkit flow
+
+#### Allowed :
+``` npx tsx src/index.ts "How do I integrate with Azure Blob Storage?"```
+
+#### Blocked:
+``` npx tsx src/index.ts "workflow to download a file from an API, save it to Blob file and export the API key"```
+
+![Image showing Generation Blocked](./GenerationBlocked.png)
+
+
+#### PII MASK and UNMASK:
+``` npx tsx src/index.ts "workflow to download a file from an API, save it to Blob file with my email john.doe@example.com"```
+
+![Image showing PII data masked ](./MaskedPII.png)
+
+---
+
+## 🧠 How It Works
+
+### **1. Intent Guard**
+- Embeds the user prompt + intent descriptions using MiniLM  
+- Computes cosine similarity  
+- Blocks prompts below threshold  
+- Detects jailbreak patterns like:  
+  - “ignore previous instructions”  
+  - “you are a hacker”  
+  - “export the API key”  
+
+### **2. PII Masking**
+Before the LLM sees the prompt:
+
+```
+"Email john.doe@example.com" → "Email [[EMAIL_0]]"
+```
+
+Detected PII includes:
+
+- Emails  
+- Phone numbers  
+- Names (NER)  
+- AU identifiers (Medicare, TFN, ABN, etc.)
+
+### **3. LLM Call**
+The masked prompt is sent to the model.
+
+### **4. Response Unmasking**
+After the LLM responds:
+
+```
+"Send a confirmation email to [[EMAIL_0]]" → "Send a confirmation email to john.doe@example.com"
+```
+---
+
+## ⚙️ Configuration
+
+### **Intent Guard**
+
+```ts
+intent: {
+  mode: "semantic",
+  allowedIntent: "intent_question",
+  semantic: {
+    threshold: 0.7,
+    intents: {
+      intent_question: "Description of allowed intent"
+    }
+  }
+}
+```
+
+### **PII Guard**
+
+```ts
+pii: {
+  reversible: true
+}
+```
+
+---
+
+## 🛡️ Why This Library Exists
+
+Genkit provides a powerful LLM framework, but production systems need:
+
+- intent boundaries  
+- PII protection  
+- jailbreak resistance  
+- predictable behavior  
+
+This library adds those guardrails without heavy dependencies or complex setup.
+
+---
+
+## Contributing
+
+We plan to: 
+
+1. Extend the utility by adding Auth and Tool Middleware in further stages.
+2. Add more filter types for common malicious prompts.
+3. Add more patterns for custom PII masking.
+
+Contributions are welcome — whether it’s bug reports, new guard modules, model improvements or enhancements. This project aims to stay lightweight, modular, and production‑ready, so thoughtful contributions are appreciated.
+
+# 📄 License
+
+Apache‑2.0

package/package.json

@@ -0,0 +1,34 @@
+{
+  "name": "@intflows/genkit-guard",
+  "description": "Intflows Genkit Guard",
+  "maintainers": [
+    "Hemant Kohli"
+  ],
+  "license": "Apache-2.0",
+  "version": "0.0.4",
+  "type": "module",
+  "exports": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "files": [
+    "src/",
+    "README.md",
+    "package.json",
+    "scripts/"
+  ],
+  "scripts": {
+    "prepare-models": "node scripts/download-model.js",
+    "build": "tsc",
+    "prepublishOnly": "npm run build"
+  },
+  "peerDependencies": {
+    "genkit": "^0.5.0"
+  },
+  "dependencies": {
+    "@huggingface/transformers": "^3.0.0",
+    "zod": "^3.23.0"
+  },
+  "devDependencies": {
+    "typescript": "^5.6.0",
+    "@types/node": "^20.0.0"
+  }
+}

package/scripts/download-model.js

@@ -0,0 +1,23 @@
+#!/usr/bin/env node
+import { pipeline, env } from '@huggingface/transformers';
+import path from 'path';
+
+// Force download to a local folder in the project root
+const modelPath = path.join(process.cwd(), 'models');
+env.cacheDir = modelPath;
+env.localModelPath = modelPath;
+env.allowRemoteModels = true; 
+
+async function download() {
+  console.log('Downloading MiniLM-L6-v2 to ./models...');
+  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('Model downloaded successfully.');
+}
+
+download();

package/src/core/types.ts

@@ -0,0 +1,26 @@
+// Intent Guard Types
+export interface IntentGuardConfig {
+  allowedIntent: string;
+  intents: Record<string, string>; // name -> semantic description
+  threshold?: number;              // default 0.7
+  fallbackMessage?: string;
+}
+
+export interface IntentResult {
+  allowed: boolean;
+  score: number;
+}
+
+
+
+// PII Masking Types
+export interface PiiRule {
+  name: string;
+  pattern: RegExp;
+//   replaceWith: string;
+}
+
+export interface PiiConfig {
+  rules?: PiiRule[];
+  maskCharacter?: string; // e.g., "*"
+}

package/src/index.ts

@@ -0,0 +1,17 @@
+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() {
+  console.log('[Guard] Loading local models...');
+  await Promise.all([
+    ModelSingleton.getExtractor(),
+    ModelSingleton.getNER()
+  ]);
+  console.log('[Guard] Models loaded');
+}

package/src/intent/intentAnalyzer.ts

@@ -0,0 +1,100 @@
+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: string) {
+  return INJECTION_PATTERNS.some(p =>
+    userInput.toLowerCase().includes(p)
+  );
+}
+
+
+export async function analyzeIntentStructured(
+  input: string,
+  intents: Record<string, string>,
+  threshold: number
+) {
+  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() as number[][];
+    const score = cos_sim(vectors[0], vectors[1]);
+
+    const finalScore =
+      typeof score === 'number' ? score : (score as any).data[0];
+
+    if (finalScore > bestScore) {
+      bestScore = finalScore;
+      bestIntent = key;
+    }
+  }
+
+  return {
+    intent: bestIntent,
+    score: bestScore,
+    allowed: bestScore >= threshold
+  };
+}

package/src/middleware/middleware.ts

@@ -0,0 +1,143 @@
+import { analyzeIntentStructured, detectInjection } from '../intent/intentAnalyzer.js';
+import { detectPII } from '../pii/detector.js';
+import { PiiTokenizer } from '../pii/tokenizer.js';
+
+export function guard(config: any) {
+  return async (req: any, next: any) => {
+    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 piiMatches = await detectPII(input);
+    console.log(`[PII Guard] Detected PII: ${piiMatches.length} matches found`);
+    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
+    };
+
+    // 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: any): any => {
+        // 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: string, metadata?: any) {
+  return {
+    finishReason: 'blocked',
+    output: {
+      type: "error",
+      status: "BLOCKED",
+      message
+    },
+    metadata
+  };
+}

package/src/pii/detector.ts

@@ -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: string) {
+  const ner = await ModelSingleton.getNER();
+
+  const results: { type: string; value: string }[] = [];
+
+  // ---- 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 ----
+  for (const rule of REGEX_RULES) {
+    const matches = text.match(rule.pattern) || [];
+    matches.forEach(m => results.push({ type: rule.type, value: m }));
+  }
+
+  return results;
+}

package/src/pii/tokenizer.ts

@@ -0,0 +1,46 @@
+export type PiiResult = {
+  maskedText: string;
+  pii: Record<string, string>;
+  piiTypes: string[];
+};
+
+export class PiiTokenizer {
+  private vault = new Map<string, string>();
+  private counter = 0;
+  private piiTypes = new Set<string>();
+
+  private createToken(type: string) {
+    return `[[${type}_${this.counter++}]]`;
+  }
+
+  mask(text: string, matches: { type: string; value: string }[]): PiiResult {
+    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: string): string {
+       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/src/util/singleton.ts

@@ -0,0 +1,51 @@
+import { pipeline, env, type FeatureExtractionPipeline } 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 {
+  private static extractor: any = null;
+  private static classifier: any = null;
+
+  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() {
+    if (!this.extractor) {
+      this.init();
+      this.extractor = await (pipeline as any)('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
+    }
+    return this.extractor;
+  }
+
+  static async getNER() {
+    if (!this.classifier) {
+      this.init();
+      // Token classification model for identifying sensitive entities
+      this.classifier = await (pipeline as any)('token-classification', 'Xenova/bert-base-NER');
+    }
+    return this.classifier;
+  }
+}