From: Stefan Gasser Date: Tue, 20 Jan 2026 19:07:03 +0000 (+0100) Subject: Add scan_roles config for role-based PII/secrets filtering (#52) X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=7cbf33c2dce90ecd4968c91a557e564268d38307;p=sgasser-llm-shield.git Add scan_roles config for role-based PII/secrets filtering (#52) Allows configuring which message roles to scan for PII and secrets. By default all roles are scanned (existing behavior). When scan_roles is set, only messages with matching roles are sent to Presidio. Use case: Skip scanning large system prompts to reduce API calls and avoid false positives on app-controlled content. - Add role field to TextSpan type - OpenAI extractor populates role from message - PII detector filters by scan_roles before Presidio calls - Secrets detector filters by scan_roles before detection --- diff --git a/config.example.yaml b/config.example.yaml index 6a660d1..f7c5faa 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -85,6 +85,16 @@ pii_detection: # - MEDICAL_LICENSE # - URL + # Which message roles to scan for PII (optional) + # By default, all roles are scanned. Set this to scan only user-controlled content: + # - user: User messages (primary source of PII) + # - tool: Tool/function call results (may contain user data) + # - function: Legacy function results (OpenAI) + # scan_roles: + # - user + # - tool + # - function + # Secrets Detection settings (Secrets Shield) # Detects private keys, API keys, tokens and other secret credentials in requests secrets_detection: @@ -138,6 +148,16 @@ secrets_detection: # Even if logging.log_content is true, secret content is never logged log_detected_types: true + # Which message roles to scan for secrets (optional) + # By default, all roles are scanned. Set this to scan only user-controlled content: + # - user: User messages (primary source of secrets) + # - tool: Tool/function call results (may contain secrets) + # - function: Legacy function results (OpenAI) + # scan_roles: + # - user + # - tool + # - function + # Logging settings logging: # SQLite database for request logs diff --git a/src/config.ts b/src/config.ts index 6b338c7..f22f698 100644 --- a/src/config.ts +++ b/src/config.ts @@ -54,6 +54,7 @@ const PIIDetectionSchema = z.object({ "IP_ADDRESS", "LOCATION", ]), + scan_roles: z.array(z.string()).optional(), }); const ServerSchema = z.object({ @@ -98,6 +99,7 @@ const SecretsDetectionSchema = z.object({ entities: z.array(z.enum(SecretEntityTypes)).default(["OPENSSH_PRIVATE_KEY", "PEM_PRIVATE_KEY"]), max_scan_chars: z.coerce.number().int().min(0).default(200000), log_detected_types: z.boolean().default(true), + scan_roles: z.array(z.string()).optional(), }); const ConfigSchema = z diff --git a/src/masking/extractors/openai.test.ts b/src/masking/extractors/openai.test.ts index 5640a46..9ae4eed 100644 --- a/src/masking/extractors/openai.test.ts +++ b/src/masking/extractors/openai.test.ts @@ -24,12 +24,14 @@ describe("OpenAI Text Extractor", () => { path: "messages[0].content", messageIndex: 0, partIndex: 0, + role: "system", }); expect(spans[1]).toEqual({ text: "Hello world", path: "messages[1].content", messageIndex: 1, partIndex: 0, + role: "user", }); }); @@ -53,12 +55,14 @@ describe("OpenAI Text Extractor", () => { path: "messages[0].content[0].text", messageIndex: 0, partIndex: 0, + role: "user", }); expect(spans[1]).toEqual({ text: "Be detailed", path: "messages[0].content[2].text", messageIndex: 0, partIndex: 2, + role: "user", }); }); @@ -76,8 +80,11 @@ describe("OpenAI Text Extractor", () => { expect(spans).toHaveLength(3); expect(spans[0].messageIndex).toBe(0); + expect(spans[0].role).toBe("system"); expect(spans[1].messageIndex).toBe(1); + expect(spans[1].role).toBe("user"); expect(spans[2].messageIndex).toBe(2); + expect(spans[2].role).toBe("assistant"); }); test("skips null/undefined content", () => { diff --git a/src/masking/extractors/openai.ts b/src/masking/extractors/openai.ts index 4cca57e..b03212b 100644 --- a/src/masking/extractors/openai.ts +++ b/src/masking/extractors/openai.ts @@ -33,6 +33,7 @@ export const openaiExtractor: RequestExtractor = path: `messages[${msgIdx}].content`, messageIndex: msgIdx, partIndex: 0, + role: msg.role, }); continue; } @@ -46,6 +47,7 @@ export const openaiExtractor: RequestExtractor = path: `messages[${msgIdx}].content[${partIdx}].text`, messageIndex: msgIdx, partIndex: partIdx, + role: msg.role, }); } } diff --git a/src/masking/types.ts b/src/masking/types.ts index e63c84d..a68869a 100644 --- a/src/masking/types.ts +++ b/src/masking/types.ts @@ -10,6 +10,7 @@ export interface TextSpan { messageIndex: number; partIndex: number; nestedPartIndex?: number; + role?: string; } export interface MaskedSpan { diff --git a/src/pii/detect.ts b/src/pii/detect.ts index ecd9bff..eb07a44 100644 --- a/src/pii/detect.ts +++ b/src/pii/detect.ts @@ -100,8 +100,15 @@ export class PIIDetector { : { language: config.pii_detection.fallback_language, usedFallback: true }; // Detect PII for each span independently + const scanRoles = config.pii_detection.scan_roles + ? new Set(config.pii_detection.scan_roles) + : null; + const spanEntities: PIIEntity[][] = await Promise.all( spans.map(async (span) => { + if (scanRoles && span.role && !scanRoles.has(span.role)) { + return []; + } if (!span.text) return []; return this.detectPII(span.text, langResult.language); }), diff --git a/src/secrets/detect.ts b/src/secrets/detect.ts index 35bfc05..2f2a760 100644 --- a/src/secrets/detect.ts +++ b/src/secrets/detect.ts @@ -95,8 +95,13 @@ export function detectSecretsInSpans( } // Detect secrets in each span + const scanRoles = config.scan_roles ? new Set(config.scan_roles) : null; + const matchCounts = new Map(); const spanLocations: SecretLocation[][] = spans.map((span) => { + if (scanRoles && span.role && !scanRoles.has(span.role)) { + return []; + } const result = detectSecrets(span.text, config); for (const match of result.matches) { matchCounts.set(match.type, (matchCounts.get(match.type) || 0) + match.count);