diff --git a/candle-binding/src/ffi/classify.rs b/candle-binding/src/ffi/classify.rs index 91e38baee..1854e15e4 100644 --- a/candle-binding/src/ffi/classify.rs +++ b/candle-binding/src/ffi/classify.rs @@ -9,6 +9,7 @@ use crate::ffi::memory::{ allocate_lora_intent_array, allocate_lora_pii_array, allocate_lora_security_array, allocate_modernbert_token_entity_array, }; +use crate::ffi::types::BertTokenEntity; use crate::ffi::types::*; use crate::model_architectures::traditional::bert::{ TRADITIONAL_BERT_CLASSIFIER, TRADITIONAL_BERT_TOKEN_CLASSIFIER, @@ -654,26 +655,66 @@ pub extern "C" fn classify_candle_bert_tokens( let entities_ptr = unsafe { allocate_bert_token_entity_array(&token_entities) }; - BertTokenClassificationResult { + return BertTokenClassificationResult { entities: entities_ptr, num_entities: token_entities.len() as i32, - } + }; } Err(e) => { println!("Candle BERT token classification failed: {}", e); - BertTokenClassificationResult { + return BertTokenClassificationResult { entities: std::ptr::null_mut(), num_entities: 0, - } + }; } } - } else { - println!("TraditionalBertTokenClassifier not initialized - call init function first"); - BertTokenClassificationResult { - entities: std::ptr::null_mut(), - num_entities: 0, + } + + // Fallback to ModernBERT token classifier (for PII detection with ModernBERT models) + if let Some(classifier) = TRADITIONAL_MODERNBERT_TOKEN_CLASSIFIER.get() { + let classifier = classifier.clone(); + match classifier.classify_tokens(text) { + Ok(token_results) => { + // Filter non-background classes; Go layer applies confidence threshold + // Keep real positions (start, end) for accurate entity extraction + let token_entities: Vec<(String, String, f32, usize, usize)> = token_results + .iter() + .filter(|(_, class_idx, _, _, _)| *class_idx > 0) + .map(|(token, class_idx, confidence, start, end)| { + ( + token.clone(), + format!("class_{}", class_idx), + *confidence, + *start, + *end, + ) + }) + .collect(); + + let entities_ptr = + unsafe { allocate_modernbert_token_entity_array(&token_entities) }; + + return BertTokenClassificationResult { + entities: entities_ptr as *mut BertTokenEntity, + num_entities: token_entities.len() as i32, + }; + } + Err(e) => { + println!("ModernBERT token classification failed: {}", e); + return BertTokenClassificationResult { + entities: std::ptr::null_mut(), + num_entities: 0, + }; + } } } + + // No classifier available + println!("No token classifier initialized (Traditional BERT, ModernBERT, or LoRA) - call init function first"); + BertTokenClassificationResult { + entities: std::ptr::null_mut(), + num_entities: 0, + } } /// Classify text using Candle BERT diff --git a/deploy/kubernetes/aibrix/semantic-router-values/values.yaml b/deploy/kubernetes/aibrix/semantic-router-values/values.yaml index ec1d43537..a6c0d281d 100644 --- a/deploy/kubernetes/aibrix/semantic-router-values/values.yaml +++ b/deploy/kubernetes/aibrix/semantic-router-values/values.yaml @@ -123,7 +123,8 @@ config: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "ORGANIZATION" # Allow - scientific terms like "photosynthesis" falsely detected as ORG - type: "system_prompt" configuration: enabled: true @@ -189,7 +190,8 @@ config: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "GPE" # Allow - country/city names in general knowledge questions - type: "semantic-cache" configuration: enabled: true diff --git a/e2e/profiles/ai-gateway/profile.go b/e2e/profiles/ai-gateway/profile.go index 8cd304064..58e070574 100644 --- a/e2e/profiles/ai-gateway/profile.go +++ b/e2e/profiles/ai-gateway/profile.go @@ -325,7 +325,7 @@ func (p *Profile) kubectlApply(ctx context.Context, kubeConfig, manifest string) } func (p *Profile) kubectlDelete(ctx context.Context, kubeConfig, manifest string) error { - return p.runKubectl(ctx, kubeConfig, "delete", "-f", manifest) + return p.runKubectl(ctx, kubeConfig, "delete", "--ignore-not-found", "-f", manifest) } func (p *Profile) runKubectl(ctx context.Context, kubeConfig string, args ...string) error { diff --git a/e2e/profiles/ai-gateway/values.yaml b/e2e/profiles/ai-gateway/values.yaml index ddb3eb930..d98a2b030 100644 --- a/e2e/profiles/ai-gateway/values.yaml +++ b/e2e/profiles/ai-gateway/values.yaml @@ -142,7 +142,8 @@ config: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "ORGANIZATION" # Allow - scientific terms like "photosynthesis" falsely detected as ORG - type: "system_prompt" configuration: enabled: true @@ -396,6 +397,10 @@ config: lora_name: general-expert use_reasoning: false plugins: + - type: "pii" + configuration: + enabled: true + pii_types_allowed: [] - type: "system_prompt" configuration: enabled: true @@ -441,7 +446,8 @@ config: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "GPE" # Allow - country/city names in general knowledge questions - type: "semantic-cache" configuration: enabled: true @@ -529,7 +535,7 @@ config: case_sensitive: false - name: "sensitive_keywords" - operator: "AND" + operator: "OR" keywords: ["SSN", "credit card"] case_sensitive: false diff --git a/e2e/profiles/aibrix/profile.go b/e2e/profiles/aibrix/profile.go index 3d6ef366a..1bf1af2cb 100644 --- a/e2e/profiles/aibrix/profile.go +++ b/e2e/profiles/aibrix/profile.go @@ -468,7 +468,7 @@ func (p *Profile) kubectlApply(ctx context.Context, kubeConfig, manifest string) } func (p *Profile) kubectlDelete(ctx context.Context, kubeConfig, manifest string) error { - return p.runKubectl(ctx, kubeConfig, "delete", "-f", manifest) + return p.runKubectl(ctx, kubeConfig, "delete", "--ignore-not-found", "-f", manifest) } func (p *Profile) runKubectl(ctx context.Context, kubeConfig string, args ...string) error { diff --git a/e2e/profiles/dynamic-config/crds/intelligentroute.yaml b/e2e/profiles/dynamic-config/crds/intelligentroute.yaml index 1c53e853f..194fecff4 100644 --- a/e2e/profiles/dynamic-config/crds/intelligentroute.yaml +++ b/e2e/profiles/dynamic-config/crds/intelligentroute.yaml @@ -100,6 +100,10 @@ spec: loraName: "general-expert" useReasoning: false plugins: + - type: "pii" + configuration: + enabled: true + pii_types_allowed: [] - type: "header_mutation" configuration: add: @@ -269,7 +273,8 @@ spec: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "ORGANIZATION" # Allow - scientific terms like "photosynthesis" falsely detected as ORG - type: "system_prompt" configuration: enabled: true @@ -503,7 +508,8 @@ spec: - type: "pii" configuration: enabled: true - pii_types_allowed: [] + pii_types_allowed: + - "GPE" # Allow - country/city names like "France" in general knowledge questions - type: "semantic-cache" configuration: enabled: true diff --git a/e2e/testcases/plugin_chain_execution.go b/e2e/testcases/plugin_chain_execution.go index df01364eb..7775236ee 100644 --- a/e2e/testcases/plugin_chain_execution.go +++ b/e2e/testcases/plugin_chain_execution.go @@ -3,6 +3,7 @@ package testcases import ( "bytes" "context" + "embed" "encoding/json" "fmt" "io" @@ -13,6 +14,9 @@ import ( "k8s.io/client-go/kubernetes" ) +//go:embed testdata/plugin_chain_cases.json +var pluginChainTestData embed.FS + func init() { pkgtestcases.Register("plugin-chain-execution", pkgtestcases.TestCase{ Description: "Test plugin chain execution order and blocking behavior", @@ -23,14 +27,29 @@ func init() { // PluginChainCase represents a test case for plugin chain execution type PluginChainCase struct { + Description string `json:"description"` Query string `json:"query"` ExpectPIIBlock bool `json:"expect_pii_block"` ExpectCacheUsed bool `json:"expect_cache_used"` ExpectPromptApplied bool `json:"expect_prompt_applied"` - Description string `json:"description"` PIITypes []string `json:"pii_types"` // Expected PII types detected } +// loadPluginChainCases loads test cases from the embedded JSON file +func loadPluginChainCases() ([]PluginChainCase, error) { + data, err := pluginChainTestData.ReadFile("testdata/plugin_chain_cases.json") + if err != nil { + return nil, fmt.Errorf("failed to read plugin chain cases: %w", err) + } + + var cases []PluginChainCase + if err := json.Unmarshal(data, &cases); err != nil { + return nil, fmt.Errorf("failed to parse plugin chain cases: %w", err) + } + + return cases, nil +} + // PluginChainResult tracks the result of a single plugin chain test type PluginChainResult struct { Query string @@ -58,38 +77,14 @@ func testPluginChainExecution(ctx context.Context, client *kubernetes.Clientset, } defer stopPortForward() - // Define test cases - testCases := []PluginChainCase{ - { - Query: "My social security number is 123-45-6789", - ExpectPIIBlock: true, - ExpectCacheUsed: false, // PII blocks before cache - ExpectPromptApplied: false, // PII blocks before prompt - Description: "PII (SSN) should block entire plugin chain", - PIITypes: []string{"US_SSN"}, - }, - { - Query: "Contact me at john.doe@example.com", - ExpectPIIBlock: true, - ExpectCacheUsed: false, - ExpectPromptApplied: false, - Description: "PII (EMAIL) should block entire plugin chain", - PIITypes: []string{"EMAIL"}, - }, - { - Query: "What is 5 + 7?", - ExpectPIIBlock: false, - ExpectCacheUsed: false, // First request, cache miss - ExpectPromptApplied: true, // Should apply math expert prompt - Description: "Clean query should pass PII and apply prompt", - }, - { - Query: "Tell me about photosynthesis", - ExpectPIIBlock: false, - ExpectCacheUsed: false, - ExpectPromptApplied: true, - Description: "Biology query should pass PII plugin", - }, + // Load test cases from JSON file + testCases, err := loadPluginChainCases() + if err != nil { + return fmt.Errorf("failed to load test cases: %w", err) + } + + if opts.Verbose { + fmt.Printf("[Test] Loaded %d test cases from testdata/plugin_chain_cases.json\n", len(testCases)) } // Run plugin chain tests @@ -181,7 +176,12 @@ func testSinglePluginChain(ctx context.Context, testCase PluginChainCase, localP // Extract plugin execution headers piiViolationHeader := resp.Header.Get("x-vsr-pii-violation") - result.PIIDetected = piiViolationHeader // Store for display purposes + piiTypesHeader := resp.Header.Get("x-vsr-pii-types") + if piiTypesHeader != "" { + result.PIIDetected = piiTypesHeader // Store detected PII types for display + } else { + result.PIIDetected = piiViolationHeader // Fallback to boolean + } result.PIIBlocked = (resp.StatusCode == http.StatusForbidden || piiViolationHeader == "true") // Check cache headers (x-vsr-cache-hit or similar) diff --git a/e2e/testcases/testdata/plugin_chain_cases.json b/e2e/testcases/testdata/plugin_chain_cases.json new file mode 100644 index 000000000..ef84d09db --- /dev/null +++ b/e2e/testcases/testdata/plugin_chain_cases.json @@ -0,0 +1,187 @@ +[ + { + "description": "PII (SSN) should block entire plugin chain", + "query": "My social security number is 123-45-6789", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["US_SSN"] + }, + { + "description": "PII (EMAIL) should block entire plugin chain", + "query": "Contact me at john.doe@example.com", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["EMAIL_ADDRESS"] + }, + { + "description": "PII (EMAIL) with name context should be blocked", + "query": "Please send the report to sarah.smith@company.org for review", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["EMAIL_ADDRESS", "PERSON"] + }, + { + "description": "PII (EMAIL) in contact request should be blocked", + "query": "You can reach me at support@techcorp.io anytime", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["EMAIL_ADDRESS"] + }, + { + "description": "PII (PHONE_NUMBER) should block plugin chain", + "query": "Call me at 555-123-4567 for more details", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["PHONE_NUMBER"] + }, + { + "description": "PII (PHONE_NUMBER) with context should be blocked", + "query": "Contact our office at 555-867-5309 for urgent matters", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["PHONE_NUMBER"] + }, + { + "description": "PII (CREDIT_CARD) should block plugin chain", + "query": "Process payment with card 4532-1234-5678-9012", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["CREDIT_CARD"] + }, + { + "description": "PII (CREDIT_CARD) Visa format should be blocked", + "query": "Use my credit card 4111111111111111 for the subscription", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["CREDIT_CARD"] + }, + { + "description": "PII (US_SSN) alternative format should be blocked", + "query": "SSN for verification: 456 78 9012", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["US_SSN"] + }, + { + "description": "PII (PERSON) with full name should be blocked", + "query": "Send the invoice to Dr. Michael Johnson at the clinic", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["PERSON"] + }, + { + "description": "PII (STREET_ADDRESS) should be blocked", + "query": "Deliver the package to 123 Main Street, Apt 4B", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["STREET_ADDRESS"] + }, + { + "description": "PII (IP_ADDRESS) IPv4 should be blocked", + "query": "The server IP is 192.168.1.100, please whitelist it", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["IP_ADDRESS"] + }, + { + "description": "PII (IP_ADDRESS) IPv6 should be blocked", + "query": "Connect to 2001:0db8:85a3:0000:0000:8a2e:0370:7334 for the service", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["IP_ADDRESS"] + }, + { + "description": "PII (IBAN_CODE) should be blocked", + "query": "Transfer funds to IBAN: DE89370400440532013000", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["IBAN_CODE"] + }, + { + "description": "Multiple PII (EMAIL + PHONE) should be blocked", + "query": "Contact support@example.com or call 555-987-6543", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["EMAIL_ADDRESS", "PHONE_NUMBER"] + }, + { + "description": "Multiple PII (SSN + PERSON) should be blocked", + "query": "Patient John Smith, SSN 234-56-7890, needs immediate care", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["US_SSN", "PERSON"] + }, + { + "description": "Multiple PII (CREDIT_CARD + EMAIL) should be blocked", + "query": "Send receipt to buyer@shop.com for card ending 4242424242424242", + "expect_pii_block": true, + "expect_cache_used": false, + "expect_prompt_applied": false, + "pii_types": ["CREDIT_CARD", "EMAIL_ADDRESS"] + }, + { + "description": "Clean math query should pass PII and apply prompt", + "query": "What is 5 + 7?", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + }, + { + "description": "Biology query should pass PII plugin (ORGANIZATION allowed)", + "query": "Tell me about photosynthesis", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + }, + { + "description": "Clean coding question should pass", + "query": "How do I write a for loop in Python?", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + }, + { + "description": "Clean general knowledge question should pass", + "query": "What is the capital of France?", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + }, + { + "description": "Clean science question should pass", + "query": "Explain quantum entanglement in simple terms", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + }, + { + "description": "Clean history question should pass", + "query": "When did World War II end?", + "expect_pii_block": false, + "expect_cache_used": false, + "expect_prompt_applied": true, + "pii_types": [] + } +] + diff --git a/src/semantic-router/pkg/headers/headers.go b/src/semantic-router/pkg/headers/headers.go index 46206ebfb..a2e1d7168 100644 --- a/src/semantic-router/pkg/headers/headers.go +++ b/src/semantic-router/pkg/headers/headers.go @@ -59,6 +59,10 @@ const ( // Value: "true" VSRPIIViolation = "x-vsr-pii-violation" + // VSRPIITypes contains the comma-separated list of PII types that were detected and denied. + // Value: "EMAIL_ADDRESS,US_SSN" (example) + VSRPIITypes = "x-vsr-pii-types" + // VSRJailbreakBlocked indicates that a jailbreak attempt was detected and blocked. // Value: "true" VSRJailbreakBlocked = "x-vsr-jailbreak-blocked" diff --git a/src/semantic-router/pkg/utils/http/response.go b/src/semantic-router/pkg/utils/http/response.go index be5c24ae5..879bff3f7 100644 --- a/src/semantic-router/pkg/utils/http/response.go +++ b/src/semantic-router/pkg/utils/http/response.go @@ -3,6 +3,7 @@ package http import ( "encoding/json" "fmt" + "strings" "time" core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" @@ -20,6 +21,9 @@ func CreatePIIViolationResponse(model string, deniedPII []string, isStreaming bo // Record PII violation metrics metrics.RecordPIIViolations(model, deniedPII) + // Join denied PII types for header + deniedPIIStr := strings.Join(deniedPII, ",") + // Create OpenAI-compatible response format for PII violations unixTimeStep := time.Now().Unix() var responseBody []byte @@ -108,7 +112,12 @@ func CreatePIIViolationResponse(model string, deniedPII []string, isStreaming bo }, }, { - // Add decision header so tests can verify the decision was made + Header: &core.HeaderValue{ + Key: headers.VSRPIITypes, + RawValue: []byte(deniedPIIStr), + }, + }, + { Header: &core.HeaderValue{ Key: headers.VSRSelectedDecision, RawValue: []byte(decisionName), diff --git a/src/semantic-router/pkg/utils/pii/policy.go b/src/semantic-router/pkg/utils/pii/policy.go index 8260c95ea..73cf63428 100644 --- a/src/semantic-router/pkg/utils/pii/policy.go +++ b/src/semantic-router/pkg/utils/pii/policy.go @@ -69,7 +69,8 @@ func (pc *PolicyChecker) CheckPolicy(decisionName string, detectedPII []string) } // If allow_by_default is false, check if this PII type is explicitly allowed - isAllowed := slices.Contains(policy.PIITypes, piiType) + // Support both exact matches and BIO-tag stripping (B-PERSON, I-PERSON → PERSON) + isAllowed := isPIITypeAllowed(piiType, policy.PIITypes) if !isAllowed { deniedPII = append(deniedPII, piiType) } @@ -84,6 +85,29 @@ func (pc *PolicyChecker) CheckPolicy(decisionName string, detectedPII []string) return true, nil, nil } +// isPIITypeAllowed checks if a PII type is in the allowed list. +// Supports exact matching and BIO-tag stripping (e.g., "B-ORGANIZATION" matches "ORGANIZATION"). +// Uses exact string matching only - no regex/wildcards. +func isPIITypeAllowed(piiType string, allowedTypes []string) bool { + // Try exact match first + if slices.Contains(allowedTypes, piiType) { + return true + } + + // Strip BIO prefix (B-, I-, O-, E-) and match base type + if len(piiType) > 2 && piiType[1] == '-' { + prefix := piiType[0] + if prefix == 'B' || prefix == 'I' || prefix == 'O' || prefix == 'E' { + baseType := piiType[2:] + if slices.Contains(allowedTypes, baseType) { + return true + } + } + } + + return false +} + // ExtractAllContent extracts all content from user and non-user messages for PII analysis func ExtractAllContent(userContent string, nonUserMessages []string) []string { var allContent []string