@intflows/genkit-guard 0.0.12 → 0.0.13

package/README.md

@@ -63,7 +63,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) 
+_This downloads the models to `./models`; the total size is approximately 1.5 GB._
 
 ### 2. Update genkit
 
@@ -193,10 +193,26 @@ intent: {
 ### **PII Guard**
 
 ```ts
-pii: {
-  reversible: true
-}
-```
+pii: {
+  reversible: true,
+  mode: "classifier"
+}
+```
+
+`classifier` mode uses `openai/privacy-filter` as a token-classification model with aggregated
+spans. Model-detected names, addresses, emails, phone numbers, URLs, dates, account numbers and
+secrets are converted into reversible masking tokens. Regex rules continue to run as an additional
+layer, and duplicate spans are masked only once.
+
+During multi-turn tool execution, opaque tokens returned through a different Genkit middleware
+context are rehydrated from the configured vault before tool execution and before the final
+response is returned to the application.
+
+Preload the same mode during application startup:
+
+```ts
+await initGuard({ pii: { mode: "classifier" } });
+```
 
 ### **PII Vault Isolation and External Storage**
 

package/dist/middleware/middleware.js

@@ -128,6 +128,11 @@ function createGuardHooks(config) {
         tool: async (req, ctx, next) => {
             const state = getGuardState(ctx);
             const toolName = req?.toolRequest?.name;
+            // Genkit may provide a fresh middleware context for a tool turn. Create a recovery
+            // tokenizer that uses the configured vault so opaque tokens can be rehydrated safely.
+            if (state.tokenizers.length === 0) {
+                state.tokenizers.push(createTokenizer(config, req, ctx));
+            }
             if (req?.toolRequest && 'input' in req.toolRequest) {
                 req.toolRequest.input = await unmaskObject(req.toolRequest.input, state.tokenizers);
             }
@@ -199,6 +204,9 @@ async function unmaskObject(obj, tokenizers) {
     return transformStrings(obj, async (value) => {
         let result = value;
         for (const tokenizer of tokenizers) {
+            // A token may have been produced by another model/tool turn with a different Genkit
+            // context or vault scope. Import only opaque tokens actually present in this value.
+            await tokenizer.importTokens(result);
             result = await tokenizer.unmask(result);
         }
         return result;

package/dist/pii/detector.d.ts

@@ -1,10 +1,20 @@
+export type PiiMatch = {
+    type: string;
+    value: string;
+};
+export type PrivacyFilterSpan = {
+    entity_group?: string;
+    entity?: string;
+    word?: string;
+    start?: number;
+    end?: number;
+    score?: number;
+};
 export declare function detectPII(text: string, opts?: {
     model?: string;
     mode?: 'ner' | 'classifier';
 }): Promise<{
-    matches: {
-        type: string;
-        value: string;
-    }[];
+    matches: PiiMatch[];
     classifier: any;
 }>;
+export declare function privacyFilterOutputToMatches(text: string, output: unknown): PiiMatch[];

package/dist/pii/detector.js

@@ -1,4 +1,14 @@
 import { ModelSingleton } from '../util/singleton.js';
+const PRIVACY_FILTER_TYPE_MAP = {
+    account_number: 'ACCOUNT_NUMBER',
+    private_address: 'ADDRESS',
+    private_email: 'EMAIL',
+    private_person: 'NAME',
+    private_phone: 'PHONE',
+    private_url: 'URL',
+    private_date: 'DATE',
+    secret: 'SECRET',
+};
 const REGEX_RULES = [
     // EMAIL (keep your existing one)
     { type: 'EMAIL', pattern: /\b[\w\.-]+@[\w\.-]+\.\w{2,}\b/gi },
@@ -36,12 +46,52 @@ export async function detectPII(text, opts) {
         }
     }
     else {
-        // classifier mode: we call the classifier and return its output alongside regex matches.
+        // Privacy Filter is a token-classification model. Aggregation produces complete spans
+        // rather than individual BIOES-labelled tokens.
         const cls = await ModelSingleton.getPIIClassifier(model);
-        classifierOutput = await cls(text);
+        classifierOutput = await cls(text, { aggregation_strategy: 'simple' });
+        for (const match of privacyFilterOutputToMatches(text, classifierOutput)) {
+            if (!results.some((existing) => existing.value === match.value)) {
+                results.push(match);
+            }
+        }
     }
     return {
         matches: results,
         classifier: classifierOutput
     };
 }
+export function privacyFilterOutputToMatches(text, output) {
+    if (!Array.isArray(output))
+        return [];
+    const matches = [];
+    for (const candidate of output) {
+        if (!candidate || typeof candidate !== 'object')
+            continue;
+        const span = candidate;
+        const rawLabel = span.entity_group ?? span.entity;
+        if (typeof rawLabel !== 'string')
+            continue;
+        const label = rawLabel.replace(/^[BIES]-/, '').toLowerCase();
+        const type = PRIVACY_FILTER_TYPE_MAP[label];
+        if (!type)
+            continue;
+        let value;
+        if (Number.isInteger(span.start) &&
+            Number.isInteger(span.end) &&
+            span.start >= 0 &&
+            span.end > span.start &&
+            span.end <= text.length) {
+            value = text.slice(span.start, span.end);
+        }
+        else if (typeof span.word === 'string') {
+            value = span.word.trim();
+        }
+        if (!value || !text.includes(value))
+            continue;
+        if (!matches.some((existing) => existing.value === value)) {
+            matches.push({ type, value });
+        }
+    }
+    return matches;
+}

package/dist/util/singleton.d.ts

@@ -1,7 +1,8 @@
+export declare const PRIVACY_FILTER_PIPELINE_TASK: "token-classification";
 export declare class ModelSingleton {
     private static extractors;
     private static nerClassifiers;
-    private static textClassifiers;
+    private static privacyFilters;
     static init(): void;
     static getExtractor(modelName?: string): Promise<any>;
     static getNER(modelName?: string): Promise<any>;

package/dist/util/singleton.js

@@ -3,12 +3,13 @@ import path from 'path';
 import fs from 'fs';
 import { fileURLToPath } from 'url';
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
+export const PRIVACY_FILTER_PIPELINE_TASK = 'token-classification';
 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 privacyFilters = new Map();
     static init() {
         // Always resolve model path relative to the client app, not the library
         const projectRoot = process.cwd();
@@ -55,12 +56,14 @@ export class ModelSingleton {
         return this.nerClassifiers.get(modelName);
     }
     static async getPIIClassifier(modelName = 'openai/privacy-filter') {
-        if (!this.textClassifiers.has(modelName)) {
+        if (!this.privacyFilters.has(modelName)) {
             this.init();
-            const inst = await pipeline('text-classification', modelName);
-            this.textClassifiers.set(modelName, inst);
+            const inst = await pipeline(PRIVACY_FILTER_PIPELINE_TASK, modelName, {
+                dtype: 'q4',
+            });
+            this.privacyFilters.set(modelName, inst);
         }
-        return this.textClassifiers.get(modelName);
+        return this.privacyFilters.get(modelName);
     }
     static async preload(models) {
         const tasks = [];

package/package.json

@@ -22,7 +22,7 @@
     "huggingface"
   ],
   "license": "Apache-2.0",
-  "version": "0.0.12",
+  "version": "0.0.13",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -34,9 +34,11 @@
   ],
   "scripts": {
     "prepare-models": "node scripts/download-model.js",
-    "build": "tsc",
-    "test": "npm run test:types && npm run test:storage && npm run test:concurrency",
+    "clean": "node -e \"require('node:fs').rmSync('dist', { recursive: true, force: true })\"",
+    "build": "npm run clean && tsc",
+    "test": "npm run test:types && npm run test:privacy-filter && npm run test:storage && npm run test:concurrency",
     "test:types": "tsc --noEmit --ignoreConfig --module NodeNext --moduleResolution NodeNext --target ESNext --strict --skipLibCheck scripts/test-types.ts",
+    "test:privacy-filter": "npm run build && node scripts/test-privacy-filter.js",
     "test:storage": "npm run build && node scripts/test-storage.js",
     "test:concurrency": "npm run build && node scripts/test-concurrent-vault.js",
     "test:redis": "npm run build && node scripts/test-redis-vault.js",

package/scripts/test-privacy-filter.js

@@ -0,0 +1,116 @@
+import assert from 'node:assert/strict';
+import { detectPII, privacyFilterOutputToMatches } from '../dist/pii/detector.js';
+import { ModelSingleton, PRIVACY_FILTER_PIPELINE_TASK } from '../dist/util/singleton.js';
+import { InMemoryPiiVaultStorage } from '../dist/pii/storage.js';
+import { PiiTokenizer } from '../dist/pii/tokenizer.js';
+import { guard } from '../dist/index.js';
+
+assert.equal(PRIVACY_FILTER_PIPELINE_TASK, 'token-classification');
+
+const input = 'Contact Alice Smith at alice@example.com. Her key is sk-live-secret.';
+const expectedSpans = [
+  { entity_group: 'private_person', word: ' Alice Smith' },
+  { entity_group: 'private_email', word: ' alice@example.com' },
+  { entity_group: 'secret', word: ' sk-live-secret' },
+];
+
+assert.deepEqual(privacyFilterOutputToMatches(input, expectedSpans), [
+  { type: 'NAME', value: 'Alice Smith' },
+  { type: 'EMAIL', value: 'alice@example.com' },
+  { type: 'SECRET', value: 'sk-live-secret' },
+]);
+
+assert.deepEqual(
+  privacyFilterOutputToMatches('Home: 10 Green Street', [
+    { entity_group: 'private_address', start: 6, end: 21, word: 'wrong fallback' },
+    { entity_group: 'O', word: 'Home' },
+    { entity_group: 'private_date', word: 'not present' },
+    null,
+  ]),
+  [{ type: 'ADDRESS', value: '10 Green Street' }]
+);
+assert.deepEqual(privacyFilterOutputToMatches(input, { invalid: true }), []);
+
+const originalGetPIIClassifier = ModelSingleton.getPIIClassifier;
+const originalGetExtractor = ModelSingleton.getExtractor;
+let receivedOptions;
+ModelSingleton.getPIIClassifier = async () => async (_text, options) => {
+  receivedOptions = options;
+  return expectedSpans;
+};
+
+try {
+  const detection = await detectPII(input, { mode: 'classifier' });
+  assert.deepEqual(receivedOptions, { aggregation_strategy: 'simple' });
+
+  // The regex email and model email are deduplicated; model-only name and secret become matches.
+  assert.deepEqual(detection.matches, [
+    { type: 'EMAIL', value: 'alice@example.com' },
+    { type: 'NAME', value: 'Alice Smith' },
+    { type: 'SECRET', value: 'sk-live-secret' },
+  ]);
+
+  const storage = new InMemoryPiiVaultStorage();
+  const tokenizer = new PiiTokenizer({ scopeId: 'privacy-filter-test', storage });
+  const masked = await tokenizer.mask(input, detection.matches);
+
+  assert.doesNotMatch(masked.maskedText, /Alice Smith|alice@example\.com|sk-live-secret/);
+  assert.match(masked.maskedText, /\[\[NAME_/);
+  assert.match(masked.maskedText, /\[\[EMAIL_/);
+  assert.match(masked.maskedText, /\[\[SECRET_/);
+  assert.equal(await tokenizer.unmask(masked.maskedText), input);
+} finally {
+  ModelSingleton.getPIIClassifier = originalGetPIIClassifier;
+}
+
+// Reproduce a Genkit multi-turn response: an inner turn creates the token in one scope, while
+// the final response passes through middleware attached to another context and scope.
+const crossTurnStorage = new InMemoryPiiVaultStorage();
+const innerTurn = new PiiTokenizer({ scopeId: 'inner-turn', storage: crossTurnStorage });
+const crossTurnMasked = await innerTurn.mask('owner@example.com', [
+  { type: 'EMAIL', value: 'owner@example.com' },
+]);
+
+ModelSingleton.getExtractor = async () => async () => ({ tolist: () => [] });
+ModelSingleton.getPIIClassifier = async () => async () => [];
+
+try {
+  const middleware = guard({
+    intent: { semantic: { threshold: 0, intents: {} } },
+    pii: {
+      mode: 'classifier',
+      vault: { storage: crossTurnStorage, scopeId: 'outer-turn' },
+    },
+    logging: { enabled: false },
+  });
+
+  const response = await middleware.model(
+    { prompt: 'Fetch blob metadata' },
+    {},
+    async () => ({ answer: `File owner: ${crossTurnMasked.maskedText}` })
+  );
+  assert.deepEqual(response, { answer: 'File owner: owner@example.com' });
+
+  const unknownTokenResponse = await middleware.model(
+    { prompt: 'Fetch blob metadata' },
+    {},
+    async () => ({ answer: '[[EMAIL_unknownnamespace_99]]' })
+  );
+  assert.deepEqual(unknownTokenResponse, { answer: '[[EMAIL_unknownnamespace_99]]' });
+
+  let toolInput;
+  await middleware.tool(
+    { toolRequest: { name: 'sendEmail', input: { recipient: crossTurnMasked.maskedText } } },
+    {},
+    async (request) => {
+      toolInput = request.toolRequest.input;
+      return { sent: true };
+    }
+  );
+  assert.deepEqual(toolInput, { recipient: 'owner@example.com' });
+} finally {
+  ModelSingleton.getExtractor = originalGetExtractor;
+  ModelSingleton.getPIIClassifier = originalGetPIIClassifier;
+}
+
+console.log('Privacy Filter pipeline, masking and cross-turn unmasking tests passed.');

package/scripts/test-types.ts

@@ -1,21 +1,21 @@
-import { guard, createRedisPiiVaultStorage, type RedisPiiVaultClient } from '../src/index.js';
-
-const redis: RedisPiiVaultClient = {
-  async hGet() {
-    return undefined;
-  },
-  async hSet() {},
-  async hGetAll() {
-    return {};
-  },
-};
-
-guard({
-  pii: {
-    reversible: true,
-    vault: {
-      storage: createRedisPiiVaultStorage(redis),
-      scopeId: (req: any, ctx: any) => ctx?.auth?.sessionId ?? req?.metadata?.requestId,
-    },
-  },
-});
+import { guard, createRedisPiiVaultStorage, type RedisPiiVaultClient } from '../src/index.js';
+
+const redis: RedisPiiVaultClient = {
+  async hGet() {
+    return undefined;
+  },
+  async hSet() {},
+  async hGetAll() {
+    return {};
+  },
+};
+
+guard({
+  pii: {
+    reversible: true,
+    vault: {
+      storage: createRedisPiiVaultStorage(redis),
+      scopeId: (req: any, ctx: any) => ctx?.auth?.sessionId ?? req?.metadata?.requestId,
+    },
+  },
+});