UNPKG

langchain

Version:
1 lines • 19.5 kB
{"version":3,"file":"piiRedaction.cjs","names":["z","HumanMessage","ToolMessage","SystemMessage","AIMessage","createMiddleware","RemoveMessage"],"sources":["../../../src/agents/middleware/piiRedaction.ts"],"sourcesContent":["import { z } from \"zod/v3\";\nimport {\n BaseMessage,\n AIMessage,\n HumanMessage,\n ToolMessage,\n RemoveMessage,\n SystemMessage,\n} from \"@langchain/core/messages\";\nimport type { InferInteropZodInput } from \"@langchain/core/utils/types\";\n\nimport { createMiddleware } from \"../middleware.js\";\n\n/**\n * Type for the redaction map that stores original values by ID\n */\ntype RedactionMap = Record<string, string>;\n\n/**\n * Configuration schema for the Input Guardrails middleware\n */\nconst contextSchema = z.object({\n /**\n * A record of PII detection rules to apply\n * @default DEFAULT_PII_RULES (with enabled rules only)\n */\n rules: z\n .record(\n z.string(),\n z.instanceof(RegExp).describe(\"Regular expression pattern to match PII\")\n )\n .optional(),\n});\n\n/**\n * @deprecated\n */\nexport type PIIRedactionMiddlewareConfig = InferInteropZodInput<\n typeof contextSchema\n>;\n\n/**\n * Generate a unique ID for a redaction\n */\nfunction generateRedactionId(): string {\n return Math.random().toString(36).substring(2, 11);\n}\n\n/**\n * Apply PII detection rules to text with ID tracking\n */\nfunction applyPIIRules(\n text: string,\n rules: Record<string, RegExp>,\n redactionMap: RedactionMap\n): string {\n let processedText = text;\n\n for (const [name, pattern] of Object.entries(rules)) {\n const replacement = name.toUpperCase().replace(/[^a-zA-Z0-9_-]/g, \"\");\n processedText = processedText.replace(pattern, (match) => {\n const id = generateRedactionId();\n redactionMap[id] = match;\n // Create a trackable replacement like [REDACTED_SSN_abc123]\n return `[REDACTED_${replacement}_${id}]`;\n });\n }\n\n return processedText;\n}\n\ninterface ProcessHumanMessageConfig {\n rules: Record<string, RegExp>;\n redactionMap: RedactionMap;\n}\n\n/**\n * Process a single human message for PII detection and redaction\n */\nasync function processMessage(\n message: BaseMessage,\n config: ProcessHumanMessageConfig\n): Promise<BaseMessage> {\n /**\n * handle basic message types\n */\n if (\n HumanMessage.isInstance(message) ||\n ToolMessage.isInstance(message) ||\n SystemMessage.isInstance(message)\n ) {\n const content = message.content as string;\n const processedContent = await applyPIIRules(\n content,\n config.rules,\n config.redactionMap\n );\n\n if (processedContent !== content) {\n const MessageConstructor = Object.getPrototypeOf(message).constructor;\n return new MessageConstructor({\n ...message,\n content: processedContent,\n });\n }\n\n return message;\n }\n\n /**\n * Handle AI messages\n */\n if (AIMessage.isInstance(message)) {\n const content =\n typeof message.content === \"string\"\n ? message.content\n : JSON.stringify(message.content);\n const toolCalls = JSON.stringify(message.tool_calls);\n const processedContent = await applyPIIRules(\n content,\n config.rules,\n config.redactionMap\n );\n const processedToolCalls = await applyPIIRules(\n toolCalls,\n config.rules,\n config.redactionMap\n );\n\n if (processedContent !== content || processedToolCalls !== toolCalls) {\n return new AIMessage({\n ...message,\n content:\n typeof message.content === \"string\"\n ? processedContent\n : JSON.parse(processedContent),\n tool_calls: JSON.parse(processedToolCalls),\n });\n }\n\n return message;\n }\n\n throw new Error(`Unsupported message type: ${message.type}`);\n}\n\n/**\n * Restore original values from redacted text using the redaction map\n */\nfunction restoreRedactedValues(\n text: string,\n redactionMap: RedactionMap\n): string {\n let restoredText = text;\n\n // Pattern to match redacted values like [REDACTED_SSN_abc123]\n const redactionPattern = /\\[REDACTED_[A-Z_]+_(\\w+)\\]/g;\n\n restoredText = restoredText.replace(redactionPattern, (match, id) => {\n if (redactionMap[id]) {\n return redactionMap[id];\n }\n return match; // Keep original if no mapping found\n });\n\n return restoredText;\n}\n\n/**\n * Restore redacted values in a message (creates a new message object)\n */\nfunction restoreMessage(\n message: BaseMessage,\n redactionMap: RedactionMap\n): { message: BaseMessage; changed: boolean } {\n /**\n * handle basic message types\n */\n if (\n HumanMessage.isInstance(message) ||\n ToolMessage.isInstance(message) ||\n SystemMessage.isInstance(message)\n ) {\n const content = message.content as string;\n const restoredContent = restoreRedactedValues(content, redactionMap);\n if (restoredContent !== content) {\n const MessageConstructor = Object.getPrototypeOf(message).constructor;\n const newMessage = new MessageConstructor({\n ...message,\n content: restoredContent,\n });\n return { message: newMessage, changed: true };\n }\n return { message, changed: false };\n }\n\n /**\n * handle AI messages\n */\n if (AIMessage.isInstance(message)) {\n const content =\n typeof message.content === \"string\"\n ? message.content\n : JSON.stringify(message.content);\n const toolCalls = JSON.stringify(message.tool_calls);\n const processedContent = restoreRedactedValues(content, redactionMap);\n const processedToolCalls = restoreRedactedValues(toolCalls, redactionMap);\n if (processedContent !== content || processedToolCalls !== toolCalls) {\n return {\n message: new AIMessage({\n ...message,\n content:\n typeof message.content === \"string\"\n ? processedContent\n : JSON.parse(processedContent),\n tool_calls: JSON.parse(processedToolCalls),\n }),\n changed: true,\n };\n }\n\n return { message, changed: false };\n }\n\n throw new Error(`Unsupported message type: ${message.type}`);\n}\n\n/**\n * Creates a middleware that detects and redacts personally identifiable information (PII)\n * from messages before they are sent to model providers, and restores original values\n * in model responses for tool execution.\n *\n * ## Mechanism\n *\n * The middleware intercepts agent execution at two points:\n *\n * ### Request Phase (`wrapModelCall`)\n * - Applies regex-based pattern matching to all message content (HumanMessage, ToolMessage, SystemMessage, AIMessage)\n * - Processes both message text and AIMessage tool call arguments\n * - Each matched pattern generates:\n * - Unique identifier: `generateRedactionId()` → `\"abc123\"`\n * - Redaction marker: `[REDACTED_{RULE_NAME}_{ID}]` → `\"[REDACTED_SSN_abc123]\"`\n * - Redaction map entry: `{ \"abc123\": \"123-45-6789\" }`\n * - Returns modified request with redacted message content\n *\n * ### Response Phase (`afterModel`)\n * - Scans AIMessage responses for redaction markers matching pattern: `/\\[REDACTED_[A-Z_]+_(\\w+)\\]/g`\n * - Replaces markers with original values from redaction map\n * - Handles both standard responses and structured output (via tool calls or JSON content)\n * - For structured output, restores values in both the tool call arguments and the `structuredResponse` state field\n * - Returns new message instances via RemoveMessage/AIMessage to update state\n *\n * ## Data Flow\n *\n * ```\n * User Input: \"My SSN is 123-45-6789\"\n * ↓ [beforeModel]\n * Model Request: \"My SSN is [REDACTED_SSN_abc123]\"\n * ↓ [model invocation]\n * Model Response: tool_call({ \"ssn\": \"[REDACTED_SSN_abc123]\" })\n * ↓ [afterModel]\n * Tool Execution: tool({ \"ssn\": \"123-45-6789\" })\n * ```\n *\n * ## Limitations\n *\n * This middleware provides model provider isolation only. PII may still be present in:\n * - LangGraph state checkpoints (memory, databases)\n * - Network traffic between client and application server\n * - Application logs and trace data\n * - Tool execution arguments and responses\n * - Final agent output\n *\n * For comprehensive PII protection, implement additional controls at the application,\n * network, and storage layers.\n *\n * @param options - Configuration options\n * @param options.rules - Record of detection rules mapping rule names to regex patterns.\n * Rule names are normalized to uppercase and used in redaction markers.\n * Patterns must use the global flag (`/pattern/g`) to match all occurrences.\n *\n * @returns Middleware instance for use with `createAgent`\n *\n * @example Basic usage with custom rules\n * ```typescript\n * import { piiRedactionMiddleware } from \"langchain\";\n * import { createAgent } from \"langchain\";\n * import { tool } from \"@langchain/core/tools\";\n * import { z } from \"zod/v3\";\n *\n * const PII_RULES = {\n * ssn: /\\b\\d{3}-?\\d{2}-?\\d{4}\\b/g,\n * email: /\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}\\b/g,\n * phone: /\\b\\d{3}[-.]?\\d{3}[-.]?\\d{4}\\b/g,\n * };\n *\n * const lookupUser = tool(async ({ ssn }) => {\n * // Receives original value: \"123-45-6789\"\n * return { name: \"John Doe\", account: \"active\" };\n * }, {\n * name: \"lookup_user\",\n * description: \"Look up user by SSN\",\n * schema: z.object({ ssn: z.string() })\n * });\n *\n * const agent = createAgent({\n * model: new ChatOpenAI({ model: \"gpt-4\" }),\n * tools: [lookupUser],\n * middleware: [piiRedactionMiddleware({ rules: PII_RULES })]\n * });\n *\n * const result = await agent.invoke({\n * messages: [new HumanMessage(\"Look up SSN 123-45-6789\")]\n * });\n * // Model request: \"Look up SSN [REDACTED_SSN_abc123]\"\n * // Model response: tool_call({ \"ssn\": \"[REDACTED_SSN_abc123]\" })\n * // Tool receives: { \"ssn\": \"123-45-6789\" }\n * ```\n *\n * @example Runtime rule configuration via context\n * ```typescript\n * const agent = createAgent({\n * model: new ChatOpenAI({ model: \"gpt-4\" }),\n * tools: [someTool],\n * middleware: [piiRedactionMiddleware()]\n * });\n *\n * // Configure rules at runtime via middleware context\n * const result = await agent.invoke(\n * { messages: [new HumanMessage(\"...\")] },\n * {\n * configurable: {\n * PIIRedactionMiddleware: {\n * rules: {\n * ssn: /\\b\\d{3}-?\\d{2}-?\\d{4}\\b/g,\n * email: /\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}\\b/g,\n * }\n * }\n * }\n * }\n * );\n * ```\n *\n * @example Custom rule patterns\n * ```typescript\n * const customRules = {\n * employee_id: /EMP-\\d{6}/g,\n * api_key: /sk-[a-zA-Z0-9]{32}/g,\n * credit_card: /\\b\\d{4}[- ]?\\d{4}[- ]?\\d{4}[- ]?\\d{4}\\b/g,\n * };\n *\n * const middleware = piiRedactionMiddleware({ rules: customRules });\n * // Generates markers like: [REDACTED_EMPLOYEE_ID_xyz789]\n * ```\n *\n * @deprecated\n */\nexport function piiRedactionMiddleware(\n options: PIIRedactionMiddlewareConfig = {}\n) {\n const redactionMap: RedactionMap = {};\n\n console.warn(\n \"DEPRECATED: piiRedactionMiddleware is deprecated. Please use piiMiddleware instead, go to https://docs.langchain.com/oss/javascript/langchain/middleware/built-in#pii-detection for more information.\"\n );\n\n return createMiddleware({\n name: \"PIIRedactionMiddleware\",\n contextSchema,\n wrapModelCall: async (request, handler) => {\n /**\n * Merge options with context, following bigTool.ts pattern\n */\n const rules = request.runtime.context.rules ?? options.rules ?? {};\n\n /**\n * If no rules are provided, skip processing\n */\n if (Object.keys(rules).length === 0) {\n return handler(request);\n }\n\n const processedMessages = await Promise.all(\n request.state.messages.map((message: BaseMessage) =>\n processMessage(message, {\n rules,\n redactionMap,\n })\n )\n );\n\n return handler({\n ...request,\n messages: processedMessages,\n });\n },\n afterModel: async (state) => {\n /**\n * If no redactions were made, skip processing\n */\n if (Object.keys(redactionMap).length === 0) {\n return;\n }\n\n const lastMessage = state.messages.at(-1);\n if (!AIMessage.isInstance(lastMessage)) {\n return;\n }\n\n /**\n * In cases where we do structured output via tool calls, we also have to look at the second last message\n * as we add a custom last message to the messages array.\n */\n const secondLastMessage = state.messages.at(-2);\n\n const { message: restoredLastMessage, changed } = restoreMessage(\n lastMessage,\n redactionMap\n );\n\n if (!changed) {\n return;\n }\n\n /**\n * Identify if the last message is a structured response and restore the values if so\n */\n let structuredResponse: Record<string, unknown> | undefined;\n if (\n AIMessage.isInstance(lastMessage) &&\n lastMessage?.tool_calls?.length === 0 &&\n typeof lastMessage.content === \"string\" &&\n lastMessage.content.startsWith(\"{\") &&\n lastMessage.content.endsWith(\"}\")\n ) {\n try {\n structuredResponse = JSON.parse(\n restoreRedactedValues(lastMessage.content, redactionMap)\n );\n } catch {\n // ignore\n }\n }\n\n /**\n * Check if the second last message is a structured response tool call\n */\n const isStructuredResponseToolCall =\n AIMessage.isInstance(secondLastMessage) &&\n secondLastMessage?.tool_calls?.length !== 0 &&\n secondLastMessage?.tool_calls?.some((call) =>\n call.name.startsWith(\"extract-\")\n );\n if (isStructuredResponseToolCall) {\n const {\n message: restoredSecondLastMessage,\n changed: changedSecondLastMessage,\n } = restoreMessage(secondLastMessage, redactionMap);\n const structuredResponseRedacted = secondLastMessage.tool_calls?.find(\n (call) => call.name.startsWith(\"extract-\")\n )?.args;\n const structuredResponse = structuredResponseRedacted\n ? JSON.parse(\n restoreRedactedValues(\n JSON.stringify(structuredResponseRedacted),\n redactionMap\n )\n )\n : undefined;\n if (changed || changedSecondLastMessage) {\n return {\n ...state,\n ...(structuredResponse ? { structuredResponse } : {}),\n messages: [\n new RemoveMessage({ id: secondLastMessage.id as string }),\n new RemoveMessage({ id: lastMessage.id as string }),\n restoredSecondLastMessage,\n restoredLastMessage,\n ],\n };\n }\n }\n\n return {\n ...state,\n ...(structuredResponse ? { structuredResponse } : {}),\n messages: [\n new RemoveMessage({ id: lastMessage.id as string }),\n restoredLastMessage,\n ],\n };\n },\n });\n}\n"],"mappings":";;;;;;;AAqBA,MAAM,gBAAgBA,OAAAA,EAAE,OAAO;;;;;AAK7B,OAAOA,OAAAA,EACJ,OACCA,OAAAA,EAAE,OAAO,GACTA,OAAAA,EAAE,WAAW,MAAM,CAAC,CAAC,SAAS,yCAAyC,CACzE,CAAC,CACA,SAAS,EACd,CAAC;;;;AAYD,SAAS,sBAA8B;CACrC,OAAO,KAAK,OAAO,CAAC,CAAC,SAAS,EAAE,CAAC,CAAC,UAAU,GAAG,EAAE;AACnD;;;;AAKA,SAAS,cACP,MACA,OACA,cACQ;CACR,IAAI,gBAAgB;CAEpB,KAAK,MAAM,CAAC,MAAM,YAAY,OAAO,QAAQ,KAAK,GAAG;EACnD,MAAM,cAAc,KAAK,YAAY,CAAC,CAAC,QAAQ,mBAAmB,EAAE;EACpE,gBAAgB,cAAc,QAAQ,UAAU,UAAU;GACxD,MAAM,KAAK,oBAAoB;GAC/B,aAAa,MAAM;GAEnB,OAAO,aAAa,YAAY,GAAG,GAAG;EACxC,CAAC;CACH;CAEA,OAAO;AACT;;;;AAUA,eAAe,eACb,SACA,QACsB;;;;CAItB,IACEC,yBAAAA,aAAa,WAAW,OAAO,KAC/BC,yBAAAA,YAAY,WAAW,OAAO,KAC9BC,yBAAAA,cAAc,WAAW,OAAO,GAChC;EACA,MAAM,UAAU,QAAQ;EACxB,MAAM,mBAAmB,MAAM,cAC7B,SACA,OAAO,OACP,OAAO,YACT;EAEA,IAAI,qBAAqB,SAAS;GAChC,MAAM,qBAAqB,OAAO,eAAe,OAAO,CAAC,CAAC;GAC1D,OAAO,IAAI,mBAAmB;IAC5B,GAAG;IACH,SAAS;GACX,CAAC;EACH;EAEA,OAAO;CACT;;;;CAKA,IAAIC,yBAAAA,UAAU,WAAW,OAAO,GAAG;EACjC,MAAM,UACJ,OAAO,QAAQ,YAAY,WACvB,QAAQ,UACR,KAAK,UAAU,QAAQ,OAAO;EACpC,MAAM,YAAY,KAAK,UAAU,QAAQ,UAAU;EACnD,MAAM,mBAAmB,MAAM,cAC7B,SACA,OAAO,OACP,OAAO,YACT;EACA,MAAM,qBAAqB,MAAM,cAC/B,WACA,OAAO,OACP,OAAO,YACT;EAEA,IAAI,qBAAqB,WAAW,uBAAuB,WACzD,OAAO,IAAIA,yBAAAA,UAAU;GACnB,GAAG;GACH,SACE,OAAO,QAAQ,YAAY,WACvB,mBACA,KAAK,MAAM,gBAAgB;GACjC,YAAY,KAAK,MAAM,kBAAkB;EAC3C,CAAC;EAGH,OAAO;CACT;CAEA,MAAM,IAAI,MAAM,6BAA6B,QAAQ,MAAM;AAC7D;;;;AAKA,SAAS,sBACP,MACA,cACQ;CACR,IAAI,eAAe;CAKnB,eAAe,aAAa,QAAQ,gCAAmB,OAAO,OAAO;EACnE,IAAI,aAAa,KACf,OAAO,aAAa;EAEtB,OAAO;CACT,CAAC;CAED,OAAO;AACT;;;;AAKA,SAAS,eACP,SACA,cAC4C;;;;CAI5C,IACEH,yBAAAA,aAAa,WAAW,OAAO,KAC/BC,yBAAAA,YAAY,WAAW,OAAO,KAC9BC,yBAAAA,cAAc,WAAW,OAAO,GAChC;EACA,MAAM,UAAU,QAAQ;EACxB,MAAM,kBAAkB,sBAAsB,SAAS,YAAY;EACnE,IAAI,oBAAoB,SAAS;GAC/B,MAAM,qBAAqB,OAAO,eAAe,OAAO,CAAC,CAAC;GAK1D,OAAO;IAAE,SAAS,IAJK,mBAAmB;KACxC,GAAG;KACH,SAAS;IACX,CAC2B;IAAG,SAAS;GAAK;EAC9C;EACA,OAAO;GAAE;GAAS,SAAS;EAAM;CACnC;;;;CAKA,IAAIC,yBAAAA,UAAU,WAAW,OAAO,GAAG;EACjC,MAAM,UACJ,OAAO,QAAQ,YAAY,WACvB,QAAQ,UACR,KAAK,UAAU,QAAQ,OAAO;EACpC,MAAM,YAAY,KAAK,UAAU,QAAQ,UAAU;EACnD,MAAM,mBAAmB,sBAAsB,SAAS,YAAY;EACpE,MAAM,qBAAqB,sBAAsB,WAAW,YAAY;EACxE,IAAI,qBAAqB,WAAW,uBAAuB,WACzD,OAAO;GACL,SAAS,IAAIA,yBAAAA,UAAU;IACrB,GAAG;IACH,SACE,OAAO,QAAQ,YAAY,WACvB,mBACA,KAAK,MAAM,gBAAgB;IACjC,YAAY,KAAK,MAAM,kBAAkB;GAC3C,CAAC;GACD,SAAS;EACX;EAGF,OAAO;GAAE;GAAS,SAAS;EAAM;CACnC;CAEA,MAAM,IAAI,MAAM,6BAA6B,QAAQ,MAAM;AAC7D;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAoIA,SAAgB,uBACd,UAAwC,CAAC,GACzC;CACA,MAAM,eAA6B,CAAC;CAEpC,QAAQ,KACN,uMACF;CAEA,OAAOC,mBAAAA,iBAAiB;EACtB,MAAM;EACN;EACA,eAAe,OAAO,SAAS,YAAY;;;;GAIzC,MAAM,QAAQ,QAAQ,QAAQ,QAAQ,SAAS,QAAQ,SAAS,CAAC;;;;GAKjE,IAAI,OAAO,KAAK,KAAK,CAAC,CAAC,WAAW,GAChC,OAAO,QAAQ,OAAO;GAGxB,MAAM,oBAAoB,MAAM,QAAQ,IACtC,QAAQ,MAAM,SAAS,KAAK,YAC1B,eAAe,SAAS;IACtB;IACA;GACF,CAAC,CACH,CACF;GAEA,OAAO,QAAQ;IACb,GAAG;IACH,UAAU;GACZ,CAAC;EACH;EACA,YAAY,OAAO,UAAU;;;;GAI3B,IAAI,OAAO,KAAK,YAAY,CAAC,CAAC,WAAW,GACvC;GAGF,MAAM,cAAc,MAAM,SAAS,GAAG,EAAE;GACxC,IAAI,CAACD,yBAAAA,UAAU,WAAW,WAAW,GACnC;;;;;GAOF,MAAM,oBAAoB,MAAM,SAAS,GAAG,EAAE;GAE9C,MAAM,EAAE,SAAS,qBAAqB,YAAY,eAChD,aACA,YACF;GAEA,IAAI,CAAC,SACH;;;;GAMF,IAAI;GACJ,IACEA,yBAAAA,UAAU,WAAW,WAAW,KAChC,aAAa,YAAY,WAAW,KACpC,OAAO,YAAY,YAAY,YAC/B,YAAY,QAAQ,WAAW,GAAG,KAClC,YAAY,QAAQ,SAAS,GAAG,GAEhC,IAAI;IACF,qBAAqB,KAAK,MACxB,sBAAsB,YAAY,SAAS,YAAY,CACzD;GACF,QAAQ,CAER;GAYF,IALEA,yBAAAA,UAAU,WAAW,iBAAiB,KACtC,mBAAmB,YAAY,WAAW,KAC1C,mBAAmB,YAAY,MAAM,SACnC,KAAK,KAAK,WAAW,UAAU,CACjC,GACgC;IAChC,MAAM,EACJ,SAAS,2BACT,SAAS,6BACP,eAAe,mBAAmB,YAAY;IAClD,MAAM,6BAA6B,kBAAkB,YAAY,MAC9D,SAAS,KAAK,KAAK,WAAW,UAAU,CAC3C,CAAC,EAAE;IACH,MAAM,qBAAqB,6BACvB,KAAK,MACH,sBACE,KAAK,UAAU,0BAA0B,GACzC,YACF,CACF,IACA,KAAA;IACJ,IAAI,WAAW,0BACb,OAAO;KACL,GAAG;KACH,GAAI,qBAAqB,EAAE,mBAAmB,IAAI,CAAC;KACnD,UAAU;MACR,IAAIE,yBAAAA,cAAc,EAAE,IAAI,kBAAkB,GAAa,CAAC;MACxD,IAAIA,yBAAAA,cAAc,EAAE,IAAI,YAAY,GAAa,CAAC;MAClD;MACA;KACF;IACF;GAEJ;GAEA,OAAO;IACL,GAAG;IACH,GAAI,qBAAqB,EAAE,mBAAmB,IAAI,CAAC;IACnD,UAAU,CACR,IAAIA,yBAAAA,cAAc,EAAE,IAAI,YAAY,GAAa,CAAC,GAClD,mBACF;GACF;EACF;CACF,CAAC;AACH"}