@intflows/genkit-guard 0.0.12 → 0.0.14

package/dist/middleware/middleware.d.ts

@@ -1,6 +1,10 @@
+import { type GuardDecision, type ToolGuardConfig } from '../core/decision.js';
+import { type PiiLabelMappings } from '../guard.config.js';
 import { z } from 'genkit';
 import { type PiiVaultStorage } from '../pii/storage.js';
 export type GuardConfig = {
+    policyVersion?: string;
+    tools?: ToolGuardConfig;
     intent?: {
         mode?: string;
         allowedIntent?: string;
@@ -13,6 +17,7 @@ export type GuardConfig = {
         reversible?: boolean;
         model?: string;
         mode?: 'ner' | 'classifier';
+        labelMappings?: PiiLabelMappings;
         vault?: {
             storage?: PiiVaultStorage;
             scopeId?: string | ((req: any, ctx: any) => string | undefined);
@@ -22,6 +27,8 @@ export type GuardConfig = {
         enabled?: boolean;
         level?: LogSeverity;
         serviceName?: string;
+        /** Awaited audit callback, independent of console logging level/enabled. Failure stops execution. */
+        onDecision?: (decision: GuardDecision) => void | Promise<void>;
     };
     models?: {
         extractor?: string;
@@ -30,6 +37,20 @@ export type GuardConfig = {
 };
 type LogSeverity = 'debug' | 'info' | 'warn' | 'error';
 export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodObject<{
+    policyVersion: z.ZodOptional<z.ZodString>;
+    tools: z.ZodOptional<z.ZodObject<{
+        defaultAction: z.ZodOptional<z.ZodEnum<["allow", "block", "redact", "approval-required"]>>;
+        rules: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodEnum<["allow", "block", "redact", "approval-required"]>>>;
+        approve: z.ZodOptional<z.ZodAny>;
+    }, "strip", z.ZodTypeAny, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }>>;
     intent: z.ZodOptional<z.ZodObject<{
         mode: z.ZodOptional<z.ZodString>;
         allowedIntent: z.ZodOptional<z.ZodString>;
@@ -62,6 +83,7 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        labelMappings: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodNullable<z.ZodString>>>;
         vault: z.ZodOptional<z.ZodObject<{
             storage: z.ZodOptional<z.ZodAny>;
             scopeId: z.ZodOptional<z.ZodAny>;
@@ -73,17 +95,19 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
             scopeId?: any;
         }>>;
     }, "strip", z.ZodTypeAny, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
         } | undefined;
     }, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
@@ -93,14 +117,17 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         enabled: z.ZodOptional<z.ZodBoolean>;
         level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
         serviceName: z.ZodOptional<z.ZodString>;
+        onDecision: z.ZodOptional<z.ZodAny>;
     }, "strip", z.ZodTypeAny, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -110,6 +137,20 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         extractor?: string | undefined;
     }>>;
 }, "passthrough", z.ZodTypeAny, z.objectOutputType<{
+    policyVersion: z.ZodOptional<z.ZodString>;
+    tools: z.ZodOptional<z.ZodObject<{
+        defaultAction: z.ZodOptional<z.ZodEnum<["allow", "block", "redact", "approval-required"]>>;
+        rules: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodEnum<["allow", "block", "redact", "approval-required"]>>>;
+        approve: z.ZodOptional<z.ZodAny>;
+    }, "strip", z.ZodTypeAny, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }>>;
     intent: z.ZodOptional<z.ZodObject<{
         mode: z.ZodOptional<z.ZodString>;
         allowedIntent: z.ZodOptional<z.ZodString>;
@@ -142,6 +183,7 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        labelMappings: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodNullable<z.ZodString>>>;
         vault: z.ZodOptional<z.ZodObject<{
             storage: z.ZodOptional<z.ZodAny>;
             scopeId: z.ZodOptional<z.ZodAny>;
@@ -153,17 +195,19 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
             scopeId?: any;
         }>>;
     }, "strip", z.ZodTypeAny, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
         } | undefined;
     }, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
@@ -173,14 +217,17 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         enabled: z.ZodOptional<z.ZodBoolean>;
         level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
         serviceName: z.ZodOptional<z.ZodString>;
+        onDecision: z.ZodOptional<z.ZodAny>;
     }, "strip", z.ZodTypeAny, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -190,6 +237,20 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         extractor?: string | undefined;
     }>>;
 }, z.ZodTypeAny, "passthrough">, z.objectInputType<{
+    policyVersion: z.ZodOptional<z.ZodString>;
+    tools: z.ZodOptional<z.ZodObject<{
+        defaultAction: z.ZodOptional<z.ZodEnum<["allow", "block", "redact", "approval-required"]>>;
+        rules: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodEnum<["allow", "block", "redact", "approval-required"]>>>;
+        approve: z.ZodOptional<z.ZodAny>;
+    }, "strip", z.ZodTypeAny, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }, {
+        defaultAction?: "allow" | "block" | "redact" | "approval-required" | undefined;
+        rules?: Record<string, "allow" | "block" | "redact" | "approval-required"> | undefined;
+        approve?: any;
+    }>>;
     intent: z.ZodOptional<z.ZodObject<{
         mode: z.ZodOptional<z.ZodString>;
         allowedIntent: z.ZodOptional<z.ZodString>;
@@ -222,6 +283,7 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         reversible: z.ZodOptional<z.ZodBoolean>;
         model: z.ZodOptional<z.ZodString>;
         mode: z.ZodOptional<z.ZodEnum<["ner", "classifier"]>>;
+        labelMappings: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodNullable<z.ZodString>>>;
         vault: z.ZodOptional<z.ZodObject<{
             storage: z.ZodOptional<z.ZodAny>;
             scopeId: z.ZodOptional<z.ZodAny>;
@@ -233,17 +295,19 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
             scopeId?: any;
         }>>;
     }, "strip", z.ZodTypeAny, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
         } | undefined;
     }, {
+        model?: string | undefined;
+        labelMappings?: Record<string, string | null> | undefined;
         mode?: "ner" | "classifier" | undefined;
         reversible?: boolean | undefined;
-        model?: string | undefined;
         vault?: {
             storage?: any;
             scopeId?: any;
@@ -253,14 +317,17 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
         enabled: z.ZodOptional<z.ZodBoolean>;
         level: z.ZodOptional<z.ZodEnum<["debug", "info", "warn", "error"]>>;
         serviceName: z.ZodOptional<z.ZodString>;
+        onDecision: z.ZodOptional<z.ZodAny>;
     }, "strip", z.ZodTypeAny, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }, {
         enabled?: boolean | undefined;
         level?: "debug" | "info" | "warn" | "error" | undefined;
         serviceName?: string | undefined;
+        onDecision?: any;
     }>>;
     models: z.ZodOptional<z.ZodObject<{
         extractor: z.ZodOptional<z.ZodString>;
@@ -271,6 +338,19 @@ export declare const guardMiddleware: import("genkit").GenerateMiddleware<z.ZodO
     }>>;
 }, z.ZodTypeAny, "passthrough">>, void>;
 export declare const guardPlugin: (pluginOptions: void) => import("@genkit-ai/ai").GenkitPluginV2;
-export declare function guard(config?: GuardConfig): (req: any, ctxOrNext: any, maybeNext?: any) => Promise<any>;
+type GuardHooks = ReturnType<typeof createGuardHooks>;
+type NativeGuard = ReturnType<typeof guardMiddleware> & GuardHooks;
+type LegacyGuard = ((req: any, ctxOrNext: any, maybeNext?: any) => Promise<any>) & GuardHooks;
+export declare function guard(config: GuardConfig & {
+    tools: ToolGuardConfig;
+}): NativeGuard;
+export declare function guard(config?: GuardConfig & {
+    tools?: undefined;
+}): LegacyGuard;
+export declare function guard(config: GuardConfig): NativeGuard | LegacyGuard;
 export declare const guardAction: typeof guard;
+declare function createGuardHooks(config?: GuardConfig): {
+    model: (req: any, ctx: any, next: any) => Promise<any>;
+    tool: (req: any, ctx: any, next: any) => Promise<any>;
+};
 export {};

package/dist/middleware/middleware.js

@@ -1,10 +1,20 @@
+import { randomUUID } from 'node:crypto';
+import { GuardToolError } from '../core/decision.js';
+import { resolveGuardModels } from '../guard.config.js';
 import { generateMiddleware, z } from 'genkit';
 import { analyzeIntentStructured, detectInjection } from '../intent/intentAnalyzer.js';
 import { detectPII } from '../pii/detector.js';
 import { PiiTokenizer } from '../pii/tokenizer.js';
 import { defaultPiiVaultStorage } from '../pii/storage.js';
 const GUARD_CONTEXT_KEY = '__genkitGuard';
+const toolActionSchema = z.enum(['allow', 'block', 'redact', 'approval-required']);
 const guardConfigSchema = z.object({
+    policyVersion: z.string().optional(),
+    tools: z.object({
+        defaultAction: toolActionSchema.optional(),
+        rules: z.record(z.string(), toolActionSchema).optional(),
+        approve: z.any().optional(),
+    }).optional(),
     intent: z.object({
         mode: z.string().optional(),
         allowedIntent: z.string().optional(),
@@ -17,6 +27,7 @@ const guardConfigSchema = z.object({
         reversible: z.boolean().optional(),
         model: z.string().optional(),
         mode: z.enum(['ner', 'classifier']).optional(),
+        labelMappings: z.record(z.string(), z.string().regex(/^[A-Z_]+$/).nullable()).optional(),
         vault: z.object({
             storage: z.any().optional(),
             scopeId: z.any().optional(),
@@ -26,6 +37,7 @@ const guardConfigSchema = z.object({
         enabled: z.boolean().optional(),
         level: z.enum(['debug', 'info', 'warn', 'error']).optional(),
         serviceName: z.string().optional(),
+        onDecision: z.any().optional(),
     }).optional(),
     models: z.object({
         extractor: z.string().optional(),
@@ -40,6 +52,10 @@ export const guardPlugin = guardMiddleware.plugin;
 export function guard(config) {
     const hooks = createGuardHooks(config);
     const baseMiddleware = guardMiddleware(config);
+    // Genkit treats every function as legacy model-only middleware, ignoring tool hooks.
+    // New tool policies must use a native reference; legacy configurations stay callable.
+    if (config?.tools)
+        return Object.assign(baseMiddleware, hooks);
     const fnRunner = async (req, ctxOrNext, maybeNext) => {
         if (typeof maybeNext === 'function') {
             return hooks.model(req, ctxOrNext, maybeNext);
@@ -62,8 +78,20 @@ export function guard(config) {
 export const guardAction = guard;
 function createGuardHooks(config) {
     const logger = createLogger(config);
+    const models = resolveGuardModels(config);
+    const decide = async (start, fields) => {
+        const decision = Object.freeze({
+            schemaVersion: '1', decisionId: randomUUID(), timestamp: new Date().toISOString(),
+            policyVersion: config?.policyVersion ?? 'unversioned',
+            latencyMs: Math.max(0, performance.now() - start), ...fields,
+        });
+        logger(decision.action === 'block' || decision.action === 'approval-required' ? 'warn' : 'info', 'guard.decision', 'Guard policy decision', { decision });
+        await config?.logging?.onDecision?.(decision);
+        return decision;
+    };
     return {
         model: async (req, ctx, next) => {
+            const started = performance.now();
             const input = getInputText(req);
             logger('info', 'guard.model.start', 'Starting guard checks for model request');
             const isInjection = await detectInjection(input);
@@ -71,17 +99,25 @@ function createGuardHooks(config) {
                 logger('warn', 'guard.intent.blocked', 'Prompt injection pattern detected', {
                     reason: 'pattern_match',
                 });
+                await decide(started, { guard: 'injection', action: 'block', reasonCode: 'INJECTION_PATTERN' });
                 return block('Prompt injection detected', {
                     reason: 'pattern_match',
                 });
             }
+            await decide(started, { guard: 'injection', action: 'allow', reasonCode: 'INJECTION_CLEAR' });
+            const intentStarted = performance.now();
             logger('info', 'guard.intent.analysis.start', 'Analyzing request intent');
-            const intentResult = await analyzeIntentStructured(input, config?.intent?.semantic?.intents ?? {}, config?.intent?.semantic?.threshold ?? 0.7);
+            const intentResult = await analyzeIntentStructured(input, config?.intent?.semantic?.intents ?? {}, config?.intent?.semantic?.threshold ?? 0.7, models.extractor);
             logger('info', 'guard.intent.analysis.complete', 'Intent analysis completed', {
                 intent: intentResult.intent,
                 score: roundScore(intentResult.score),
                 allowed: intentResult.allowed,
             });
+            await decide(intentStarted, {
+                guard: 'intent', action: intentResult.allowed ? 'allow' : 'block',
+                reasonCode: intentResult.allowed ? 'INTENT_ALLOWED' : 'INTENT_REJECTED',
+                confidence: intentResult.score,
+            });
             if (!intentResult.allowed) {
                 logger('warn', 'guard.intent.blocked', 'Intent not allowed', {
                     intent: intentResult.intent,
@@ -92,6 +128,7 @@ function createGuardHooks(config) {
                     score: intentResult.score,
                 });
             }
+            const piiStarted = performance.now();
             const textForPii = collectModelRequestText(req);
             const piiResponse = await scanPII(textForPii, config);
             const piiMatches = piiResponse?.matches || [];
@@ -100,6 +137,8 @@ function createGuardHooks(config) {
             await tokenizer.importTokens(textForPii);
             await maskModelRequest(req, tokenizer, piiMatches);
             pushTokenizer(ctx, tokenizer);
+            await decide(piiStarted, { guard: 'pii', action: piiMatches.length ? 'redact' : 'allow',
+                reasonCode: piiMatches.length ? 'PII_DETECTED' : 'PII_CLEAR' });
             logger(piiMatches.length > 0 ? 'warn' : 'info', 'guard.model.pii.masked', 'PII scan completed for model request', {
                 piiDetected: piiMatches.length > 0,
                 piiMatchCount: piiMatches.length,
@@ -126,11 +165,44 @@ function createGuardHooks(config) {
             return unmaskedResponse;
         },
         tool: async (req, ctx, next) => {
-            const state = getGuardState(ctx);
+            const started = performance.now();
             const toolName = req?.toolRequest?.name;
+            const stop = async (action, reasonCode) => {
+                throw new GuardToolError(await decide(started, { guard: 'tool', action, reasonCode }));
+            };
+            const rules = config?.tools?.rules;
+            const action = rules && Object.hasOwn(rules, toolName)
+                ? rules[toolName] : config?.tools?.defaultAction ?? 'allow';
+            if (!['allow', 'block', 'redact', 'approval-required'].includes(action)) {
+                return stop('block', 'TOOL_POLICY_ERROR');
+            }
+            if (action === 'block')
+                return stop('block', 'TOOL_BLOCKED');
+            if (action === 'approval-required' && !config?.tools?.approve) {
+                return stop('approval-required', 'TOOL_APPROVAL_REQUIRED');
+            }
+            const state = getGuardState(ctx);
+            // 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);
             }
+            if (action === 'approval-required') {
+                let approved;
+                try {
+                    approved = await config.tools.approve({
+                        toolName, input: structuredClone(req?.toolRequest?.input), context: ctx?.context,
+                    });
+                }
+                catch {
+                    return stop('block', 'TOOL_POLICY_ERROR');
+                }
+                if (approved !== true)
+                    return stop('block', 'TOOL_APPROVAL_DENIED');
+            }
             const toolInputText = collectStrings(req?.toolRequest?.input).join('\n');
             const piiResponse = await scanPII(toolInputText, config);
             const piiMatches = piiResponse?.matches || [];
@@ -147,6 +219,18 @@ function createGuardHooks(config) {
                 piiMatchCount: piiMatches.length,
                 piiTypes,
             });
+            if (action === 'redact' && req?.toolRequest) {
+                req.toolRequest.input = await transformStrings(req.toolRequest.input, (value) => {
+                    // Irreversible redaction: do not send recoverable vault tokens to this tool.
+                    for (const match of [...piiMatches].sort((a, b) => b.value.length - a.value.length)) {
+                        if (match.value)
+                            value = value.split(match.value).join('[REDACTED]');
+                    }
+                    return value;
+                });
+            }
+            await decide(started, { guard: 'tool', action: action === 'redact' ? 'redact' : 'allow',
+                reasonCode: action === 'redact' ? 'TOOL_REDACTED' : action === 'approval-required' ? 'TOOL_APPROVED' : 'TOOL_ALLOWED' });
             const res = await next(req, ctx);
             const toolResponseText = collectStrings(res).join('\n');
             if (toolResponseText) {
@@ -199,6 +283,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;
@@ -241,9 +328,11 @@ async function scanPII(text, config) {
             classifier: undefined,
         };
     }
+    const models = resolveGuardModels(config);
     return detectPII(text, {
-        model: config?.pii?.model,
-        mode: config?.pii?.mode,
+        model: models.pii,
+        mode: models.mode,
+        labelMappings: config?.pii?.labelMappings,
     });
 }
 function getGuardState(ctx = {}) {

package/dist/pii/detector.d.ts

@@ -1,10 +1,22 @@
+import type { PiiLabelMappings } from '../guard.config.js';
+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';
+    labelMappings?: PiiLabelMappings;
 }): Promise<{
-    matches: {
-        type: string;
-        value: string;
-    }[];
+    matches: PiiMatch[];
     classifier: any;
 }>;
+export declare function privacyFilterOutputToMatches(text: string, output: unknown, labelMappings?: PiiLabelMappings): 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 },
@@ -30,18 +40,77 @@ export async function detectPII(text, opts) {
         const ner = await ModelSingleton.getNER(model);
         const entities = await ner(text);
         for (const e of entities) {
-            if (e.entity && e.entity.includes('PER')) {
+            const mappedType = mappedLabel(e.entity_group ?? e.entity, opts?.labelMappings);
+            if (mappedType !== undefined) {
+                if (mappedType)
+                    results.push(...privacyFilterOutputToMatches(text, [e], opts?.labelMappings));
+            }
+            else 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.
+        // 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, opts?.labelMappings)) {
+            if (!results.some((existing) => existing.value === match.value)) {
+                results.push(match);
+            }
+        }
     }
     return {
         matches: results,
         classifier: classifierOutput
     };
 }
+export function privacyFilterOutputToMatches(text, output, labelMappings) {
+    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 customType = mappedLabel(rawLabel, labelMappings);
+        const type = customType !== undefined ? customType : (Object.hasOwn(PRIVACY_FILTER_TYPE_MAP, label) ? PRIVACY_FILTER_TYPE_MAP[label] : undefined);
+        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;
+}
+function mappedLabel(label, mappings) {
+    if (typeof label !== 'string' || !mappings)
+        return undefined;
+    const normalize = (value) => value.replace(/^[BIES]-/i, '').toLowerCase();
+    const entry = Object.entries(mappings).find(([key]) => normalize(key) === normalize(label));
+    if (!entry)
+        return undefined;
+    const type = entry[1];
+    if (type !== null && !/^[A-Z_]+$/.test(type)) {
+        throw new Error('PII label mapping types must contain only uppercase letters and underscores');
+    }
+    return type;
+}

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.14",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -34,13 +34,18 @@
   ],
   "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 && npm run test:config && npm run test:release1",
     "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",
-    "prepublishOnly": "npm run build"
+    "prepublishOnly": "npm run build",
+    "test:config": "npm run build && node scripts/test-config.js",
+    "wiki:publish": "node scripts/publish-wiki.js",
+    "test:release1": "npm run build && node --test scripts/test-release1.js scripts/test-publish-wiki.js"
   },
   "dependencies": {
     "@huggingface/transformers": "^4.2.0",