@intflows/genkit-guard 0.0.13 → 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,8 +165,23 @@ 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) {
@@ -136,6 +190,19 @@ function createGuardHooks(config) {
             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 || [];
@@ -152,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) {
@@ -249,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,3 +1,4 @@
+import type { PiiLabelMappings } from '../guard.config.js';
 export type PiiMatch = {
     type: string;
     value: string;
@@ -13,8 +14,9 @@ export type PrivacyFilterSpan = {
 export declare function detectPII(text: string, opts?: {
     model?: string;
     mode?: 'ner' | 'classifier';
+    labelMappings?: PiiLabelMappings;
 }): Promise<{
     matches: PiiMatch[];
     classifier: any;
 }>;
-export declare function privacyFilterOutputToMatches(text: string, output: unknown): PiiMatch[];
+export declare function privacyFilterOutputToMatches(text: string, output: unknown, labelMappings?: PiiLabelMappings): PiiMatch[];

package/dist/pii/detector.js

@@ -40,7 +40,12 @@ 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, '') });
             }
         }
@@ -50,7 +55,7 @@ export async function detectPII(text, opts) {
         // rather than individual BIOES-labelled tokens.
         const cls = await ModelSingleton.getPIIClassifier(model);
         classifierOutput = await cls(text, { aggregation_strategy: 'simple' });
-        for (const match of privacyFilterOutputToMatches(text, classifierOutput)) {
+        for (const match of privacyFilterOutputToMatches(text, classifierOutput, opts?.labelMappings)) {
             if (!results.some((existing) => existing.value === match.value)) {
                 results.push(match);
             }
@@ -61,7 +66,7 @@ export async function detectPII(text, opts) {
         classifier: classifierOutput
     };
 }
-export function privacyFilterOutputToMatches(text, output) {
+export function privacyFilterOutputToMatches(text, output, labelMappings) {
     if (!Array.isArray(output))
         return [];
     const matches = [];
@@ -73,7 +78,8 @@ export function privacyFilterOutputToMatches(text, output) {
         if (typeof rawLabel !== 'string')
             continue;
         const label = rawLabel.replace(/^[BIES]-/, '').toLowerCase();
-        const type = PRIVACY_FILTER_TYPE_MAP[label];
+        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;
@@ -95,3 +101,16 @@ export function privacyFilterOutputToMatches(text, output) {
     }
     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/package.json

@@ -22,7 +22,7 @@
     "huggingface"
   ],
   "license": "Apache-2.0",
-  "version": "0.0.13",
+  "version": "0.0.14",
   "type": "module",
   "exports": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -36,13 +36,16 @@
     "prepare-models": "node scripts/download-model.js",
     "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": "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",

package/scripts/publish-wiki.js

@@ -0,0 +1,59 @@
+#!/usr/bin/env node
+import { spawnSync } from 'node:child_process';
+import { mkdtempSync, readdirSync, copyFileSync, rmSync, realpathSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { dirname, join, resolve } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+const root = resolve(dirname(fileURLToPath(import.meta.url)), '..');
+const args = process.argv.slice(2);
+let source = join(root, 'docs/wiki-v0.0.14');
+let repository = 'https://github.com/IntFlows/genkit-guard.wiki.git';
+let publish = false;
+for (let i = 0; i < args.length; i++) {
+  if (args[i] === '--publish') publish = true;
+  else if (['--source', '--repo'].includes(args[i])) {
+    const flag = args[i];
+    const value = args[++i];
+    if (!value || value.startsWith('-')) throw new Error(`Missing value for ${flag}`);
+    if (flag === '--source') source = resolve(value);
+    else repository = value;
+  } else if (args[i] === '--help') {
+    console.log('node scripts/publish-wiki.js [--source DIRECTORY] [--repo URL] [--publish]');
+    console.log('Default: clone wiki and preview diff. --publish commits and pushes changed numbered pages.');
+    process.exit(0);
+  } else throw new Error(`Unknown option: ${args[i]}`);
+}
+source = realpathSync(source);
+const pages = readdirSync(source, { withFileTypes: true })
+  .filter(entry => entry.isFile() && /^\d+\..+\.md$/.test(entry.name)).map(entry => entry.name).sort();
+if (!pages.length) throw new Error('No numbered Markdown wiki pages found');
+const work = mkdtempSync(join(tmpdir(), 'genkit-guard-wiki-'));
+const checkout = join(work, 'wiki');
+function git(args, cwd = work) {
+  const result = spawnSync('git', args, { cwd, encoding: 'utf8', shell: false });
+  if (result.error) throw result.error;
+  if (result.status !== 0) throw new Error(result.stderr || result.stdout || 'Git command failed');
+  return result.stdout;
+}
+try {
+  git(['clone', '--', repository, checkout]);
+  for (const page of pages) copyFileSync(join(source, page), join(checkout, page));
+  git(['add', '--', ...pages], checkout);
+  const summary = git(['diff', '--cached', '--stat'], checkout);
+  if (!summary.trim()) console.log('Wiki already matches these pages; nothing to publish.');
+  else {
+    console.log(summary);
+    console.log(git(['diff', '--cached', '--'], checkout));
+    if (!publish) console.log('Preview only. Run again with --publish to commit and push.');
+    else {
+      git(['commit', '-m', 'Update Genkit Guard wiki documentation'], checkout);
+      // Normal push: concurrent remote updates reject safely; never force-push.
+      console.log(git(['push', 'origin', 'HEAD'], checkout));
+      console.log('Wiki changes published.');
+    }
+  }
+} finally {
+  // Only remove the temporary directory created by this process.
+  rmSync(work, { recursive: true, force: true });
+}

package/scripts/test-config.js

@@ -0,0 +1,45 @@
+import assert from 'node:assert/strict';
+import { guard, initGuard, defineGuardConfig } from '../dist/index.js';
+import { ModelSingleton } from '../dist/util/singleton.js';
+import { detectPII, privacyFilterOutputToMatches } from '../dist/pii/detector.js';
+const calls = [];
+const originals = {};
+for (const method of ['getExtractor', 'getNER', 'getPIIClassifier']) {
+  originals[method] = ModelSingleton[method];
+  ModelSingleton[method] = async (name) => {
+    calls.push([method, name]);
+    return async () => method === 'getExtractor' ? { tolist: () => [[1, 0], [1, 0]] } : [{ entity: 'B-LABEL_1', word: 'Alice' }];
+  };
+}
+try {
+  await initGuard();
+  assert.deepEqual(calls.splice(0), [['getExtractor', 'Xenova/all-MiniLM-L6-v2'], ['getNER', 'Xenova/bert-base-NER']]);
+  for (const mode of ['ner', 'classifier']) {
+    const config = defineGuardConfig({
+      models: { extractor: 'custom/intent' },
+      intent: { semantic: { intents: { support: 'Customer support' } } },
+      pii: { mode, model: 'custom/pii', labelMappings: { LABEL_1: 'NAME' } },
+      logging: { enabled: false },
+    });
+    await initGuard(config);
+    const result = await guard(config).model({ prompt: 'Help Alice' }, {}, async (req) => {
+      assert.match(req.prompt, /\[\[NAME_/);
+      assert.doesNotMatch(req.prompt, /Alice/);
+      return { text: req.prompt };
+    });
+    assert.equal(result.text, 'Help Alice');
+    const expected = [['getExtractor', 'custom/intent'], [mode === 'ner' ? 'getNER' : 'getPIIClassifier', 'custom/pii']];
+    assert.deepEqual(calls.splice(0), [...expected, ...expected]);
+  }
+  await initGuard({ pii: { mode: 'classifier' } });
+  assert.equal(calls.splice(0)[1][1], 'openai/privacy-filter');
+  assert.deepEqual(privacyFilterOutputToMatches('Alice', [{ entity: 'S-private_person', word: 'Alice' }], { private_person: null }), []);
+  assert.deepEqual(privacyFilterOutputToMatches('Alice', [{ entity: 'B-label_1', word: 'Alice' }], { LABEL_1: 'CUSTOM_NAME' }), [{ type: 'CUSTOM_NAME', value: 'Alice' }]);
+  assert.deepEqual(privacyFilterOutputToMatches('Alice', [{ entity: 'toString', word: 'Alice' }]), []);
+  assert.throws(() => privacyFilterOutputToMatches('Alice', [{ entity: 'LABEL_1', word: 'Alice' }], { LABEL_1: 'bad-token' }), /uppercase/);
+  const regex = await detectPII('alice@example.com', { mode: 'classifier', labelMappings: { LABEL_1: null } });
+  assert.equal(regex.matches[0].type, 'EMAIL');
+} finally {
+  Object.assign(ModelSingleton, originals);
+}
+console.log('Shared model configuration and compatibility tests passed.');