Add banned terms guardrails

application_pages/chat_app.py

@@ -82,6 +82,13 @@ def initialize_guardrails():
                     guardrail_name,
                 )()
             )
+        elif guardrail_name == "RestrictedTermsJudge":
+            st.session_state.guardrails.append(
+                getattr(
+                    importlib.import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )()
+            )
     st.session_state.guardrails_manager = GuardrailManager(
         guardrails=st.session_state.guardrails
     )

guardrails_genie/guardrails/__init__.py

@@ -6,6 +6,7 @@ from .entity_recognition import (
     PresidioEntityRecognitionGuardrail,
     RegexEntityRecognitionGuardrail,
     TransformersEntityRecognitionGuardrail,
+    RestrictedTermsJudge,
 )
 from .manager import GuardrailManager
 
@@ -15,5 +16,6 @@ __all__ = [
     "PresidioEntityRecognitionGuardrail",
     "RegexEntityRecognitionGuardrail",
     "TransformersEntityRecognitionGuardrail",
+    "RestrictedTermsJudge",
     "GuardrailManager",
 ]

guardrails_genie/guardrails/entity_recognition/__init__.py

@@ -1,9 +1,10 @@
 from .presidio_entity_recognition_guardrail import PresidioEntityRecognitionGuardrail 
 from .regex_entity_recognition_guardrail import RegexEntityRecognitionGuardrail
 from .transformers_entity_recognition_guardrail import TransformersEntityRecognitionGuardrail
-
+from .llm_judge_entity_recognition_guardrail import RestrictedTermsJudge
 __all__ = [
     "PresidioEntityRecognitionGuardrail",
     "RegexEntityRecognitionGuardrail",
     "TransformersEntityRecognitionGuardrail",
+    "RestrictedTermsJudge"
 ]

guardrails_genie/guardrails/entity_recognition/banned_terms_examples/banned_term_examples.py

@@ -0,0 +1,178 @@
+"""
+Collection of restricted terms test examples with expected outcomes for entity recognition testing.
+Focuses on banned terms, competitor mentions, and brand protection scenarios.
+"""
+
+RESTRICTED_TERMS_EXAMPLES = [
+    {
+        "description": "Competitor Product Discussion",
+        "input_text": """
+I think we should implement features similar to Salesforce's Einstein AI 
+and Oracle's Cloud Infrastructure. Maybe we could also look at how 
+AWS handles their lambda functions.
+""",
+        "custom_terms": ["Salesforce", "Oracle", "AWS", "Einstein AI", "Cloud Infrastructure", "lambda"],
+        "expected_entities": {
+            "Salesforce": ["Salesforce"],
+            "Oracle": ["Oracle"],
+            "AWS": ["AWS"],
+            "Einstein AI": ["Einstein AI"],
+            "Cloud Infrastructure": ["Cloud Infrastructure"],
+            "lambda": ["lambda"]
+        }
+    },
+    {
+        "description": "Inappropriate Language in Support Ticket",
+        "input_text": """
+This damn product keeps crashing! What the hell is wrong with your 
+stupid service? I've wasted so much freaking time on this crap.
+""",
+        "custom_terms": ["damn", "hell", "stupid", "crap"],
+        "expected_entities": {
+            "damn": ["damn"],
+            "hell": ["hell"],
+            "stupid": ["stupid"],
+            "crap": ["crap"]
+        }
+    },
+    {
+        "description": "Confidential Project Names",
+        "input_text": """
+Project Titan's launch date has been moved up. We should coordinate 
+with Project Phoenix team and the Blue Dragon initiative for resource allocation.
+""",
+        "custom_terms": ["Project Titan", "Project Phoenix", "Blue Dragon"],
+        "expected_entities": {
+            "Project Titan": ["Project Titan"],
+            "Project Phoenix": ["Project Phoenix"],
+            "Blue Dragon": ["Blue Dragon"]
+        }
+    }
+]
+
+# Edge cases and special formats
+EDGE_CASE_EXAMPLES = [
+    {
+        "description": "Common Corporate Abbreviations and Stock Symbols",
+        "input_text": """
+MSFT's Azure and O365 platform is gaining market share. 
+Have you seen what GOOGL/GOOG and FB/META are doing with their AI?
+CRM (Salesforce) and ORCL (Oracle) have interesting features too.
+""",
+        "custom_terms": ["Microsoft", "Google", "Meta", "Facebook", "Salesforce", "Oracle"],
+        "expected_entities": {
+            "Microsoft": ["MSFT"],
+            "Google": ["GOOGL", "GOOG"],
+            "Meta": ["META"],
+            "Facebook": ["FB"],
+            "Salesforce": ["CRM", "Salesforce"],
+            "Oracle": ["ORCL"]
+        }
+    },
+    {
+        "description": "L33t Speak and Intentional Obfuscation",
+        "input_text": """
+S4l3sf0rc3 is better than 0r4cl3! 
+M1cr0$oft and G00gl3 are the main competitors.
+Let's check F8book and Met@ too.
+""",
+        "custom_terms": ["Salesforce", "Oracle", "Microsoft", "Google", "Facebook", "Meta"],
+        "expected_entities": {
+            "Salesforce": ["S4l3sf0rc3"],
+            "Oracle": ["0r4cl3"],
+            "Microsoft": ["M1cr0$oft"],
+            "Google": ["G00gl3"],
+            "Facebook": ["F8book"],
+            "Meta": ["Met@"]
+        }
+    },
+    {
+        "description": "Case Variations and Partial Matches",
+        "input_text": """
+salesFORCE and ORACLE are competitors.
+MicroSoft and google are too.
+Have you tried micro-soft or Google_Cloud?
+""",
+        "custom_terms": ["Microsoft", "Google", "Salesforce", "Oracle"],
+        "expected_entities": {
+            "Microsoft": ["MicroSoft", "micro-soft"],
+            "Google": ["google", "Google_Cloud"],
+            "Salesforce": ["salesFORCE"],
+            "Oracle": ["ORACLE"]
+        }
+    },
+    {
+        "description": "Common Misspellings and Typos",
+        "input_text": """
+Microsft and Microsooft are common typos.
+Goggle, Googel, and Gooogle are search related.
+Salezforce and Oracel need checking too.
+""",
+        "custom_terms": ["Microsoft", "Google", "Salesforce", "Oracle"],
+        "expected_entities": {
+            "Microsoft": ["Microsft", "Microsooft"],
+            "Google": ["Goggle", "Googel", "Gooogle"],
+            "Salesforce": ["Salezforce"],
+            "Oracle": ["Oracel"]
+        }
+    },
+    {
+        "description": "Mixed Variations and Context",
+        "input_text": """
+The M$ cloud competes with AWS (Amazon Web Services).
+FB/Meta's social platform and GOOGL's search dominate.
+SF.com and Oracle-DB are industry standards.
+""",
+        "custom_terms": ["Microsoft", "Amazon Web Services", "Facebook", "Meta", "Google", "Salesforce", "Oracle"],
+        "expected_entities": {
+            "Microsoft": ["M$"],
+            "Amazon Web Services": ["AWS"],
+            "Facebook": ["FB"],
+            "Meta": ["Meta"],
+            "Google": ["GOOGL"],
+            "Salesforce": ["SF.com"],
+            "Oracle": ["Oracle-DB"]
+        }
+    }
+]
+
+def validate_entities(detected: dict, expected: dict) -> bool:
+    """Compare detected entities with expected entities"""
+    if set(detected.keys()) != set(expected.keys()):
+        return False
+    return all(set(detected[k]) == set(expected[k]) for k in expected.keys())
+
+def run_test_case(guardrail, test_case, test_type="Main"):
+    """Run a single test case and print results"""
+    print(f"\n{test_type} Test Case: {test_case['description']}")
+    print("-" * 50)
+    
+    result = guardrail.guard(
+        test_case['input_text'],
+        custom_terms=test_case['custom_terms']
+    )
+    expected = test_case['expected_entities']
+    
+    # Validate results
+    matches = validate_entities(result.detected_entities, expected)
+    
+    print(f"Test Status: {'✓ PASS' if matches else '✗ FAIL'}")
+    print(f"Contains Restricted Terms: {result.contains_entities}")
+    
+    if not matches:
+        print("\nEntity Comparison:")
+        all_entity_types = set(list(result.detected_entities.keys()) + list(expected.keys()))
+        for entity_type in all_entity_types:
+            detected = set(result.detected_entities.get(entity_type, []))
+            expected_set = set(expected.get(entity_type, []))
+            print(f"\nEntity Type: {entity_type}")
+            print(f"  Expected: {sorted(expected_set)}")
+            print(f"  Detected: {sorted(detected)}")
+            if detected != expected_set:
+                print(f"  Missing: {sorted(expected_set - detected)}")
+                print(f"  Extra: {sorted(detected - expected_set)}")
+    
+    if result.anonymized_text:
+        print(f"\nAnonymized Text:\n{result.anonymized_text}")
+    
+    return matches

guardrails_genie/guardrails/entity_recognition/banned_terms_examples/run_llm_judge.py

@@ -0,0 +1,47 @@
+from guardrails_genie.guardrails.entity_recognition.llm_judge_entity_recognition_guardrail import RestrictedTermsJudge
+from guardrails_genie.guardrails.entity_recognition.banned_terms_examples.banned_term_examples import (
+    RESTRICTED_TERMS_EXAMPLES, 
+    EDGE_CASE_EXAMPLES, 
+    run_test_case
+)
+from guardrails_genie.llm import OpenAIModel
+import weave
+
+def test_restricted_terms_detection():
+    """Test restricted terms detection scenarios using predefined test cases"""
+    weave.init("guardrails-genie-restricted-terms-llm-judge")
+    
+    # Create the guardrail with OpenAI model
+    llm_judge = RestrictedTermsJudge(
+        should_anonymize=True,
+        llm_model=OpenAIModel()
+    )
+
+    # Test statistics
+    total_tests = len(RESTRICTED_TERMS_EXAMPLES) + len(EDGE_CASE_EXAMPLES)
+    passed_tests = 0
+
+    # Test main restricted terms examples
+    print("\nRunning Main Restricted Terms Tests")
+    print("=" * 80)
+    for test_case in RESTRICTED_TERMS_EXAMPLES:
+        if run_test_case(llm_judge, test_case):
+            passed_tests += 1
+
+    # Test edge cases
+    print("\nRunning Edge Cases")
+    print("=" * 80)
+    for test_case in EDGE_CASE_EXAMPLES:
+        if run_test_case(llm_judge, test_case, "Edge"):
+            passed_tests += 1
+
+    # Print summary
+    print("\nTest Summary")
+    print("=" * 80)
+    print(f"Total Tests: {total_tests}")
+    print(f"Passed: {passed_tests}")
+    print(f"Failed: {total_tests - passed_tests}")
+    print(f"Success Rate: {(passed_tests/total_tests)*100:.1f}%")
+
+if __name__ == "__main__":
+    test_restricted_terms_detection()

guardrails_genie/guardrails/entity_recognition/banned_terms_examples/run_regex_model.py

@@ -0,0 +1,46 @@
+from guardrails_genie.guardrails.entity_recognition.regex_entity_recognition_guardrail import RegexEntityRecognitionGuardrail
+from guardrails_genie.guardrails.entity_recognition.banned_terms_examples.banned_term_examples import (
+    RESTRICTED_TERMS_EXAMPLES, 
+    EDGE_CASE_EXAMPLES, 
+    run_test_case
+)
+import weave
+
+def test_restricted_terms_detection():
+    """Test restricted terms detection scenarios using predefined test cases"""
+    weave.init("guardrails-genie-restricted-terms-regex-model")
+    
+    # Create the guardrail with anonymization enabled
+    regex_guardrail = RegexEntityRecognitionGuardrail(
+        use_defaults=False,  # Don't use default PII patterns
+        should_anonymize=True
+    )
+
+    # Test statistics
+    total_tests = len(RESTRICTED_TERMS_EXAMPLES) + len(EDGE_CASE_EXAMPLES)
+    passed_tests = 0
+
+    # Test main restricted terms examples
+    print("\nRunning Main Restricted Terms Tests")
+    print("=" * 80)
+    for test_case in RESTRICTED_TERMS_EXAMPLES:
+        if run_test_case(regex_guardrail, test_case):
+            passed_tests += 1
+
+    # Test edge cases
+    print("\nRunning Edge Cases")
+    print("=" * 80)
+    for test_case in EDGE_CASE_EXAMPLES:
+        if run_test_case(regex_guardrail, test_case, "Edge"):
+            passed_tests += 1
+
+    # Print summary
+    print("\nTest Summary")
+    print("=" * 80)
+    print(f"Total Tests: {total_tests}")
+    print(f"Passed: {passed_tests}")
+    print(f"Failed: {total_tests - passed_tests}")
+    print(f"Success Rate: {(passed_tests/total_tests)*100:.1f}%")
+
+if __name__ == "__main__":
+    test_restricted_terms_detection()

guardrails_genie/guardrails/entity_recognition/llm_judge_entity_recognition_guardrail.py

@@ -1,3 +1,166 @@
-## Word conssitentcy
-# - Scent -> Odor
-# - odour -> Odor
+from typing import Dict, List, Optional
+import weave
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+from ...llm import OpenAIModel
+from ..base import Guardrail
+import instructor
+
+
+class TermMatch(BaseModel):
+    """Represents a matched term and its variations"""
+    original_term: str
+    matched_text: str
+    match_type: str = Field(
+        description="Type of match: EXACT, MISSPELLING, ABBREVIATION, or VARIANT"
+    )
+    explanation: str = Field(
+        description="Explanation of why this is considered a match"
+    )
+
+
+class RestrictedTermsAnalysis(BaseModel):
+    """Analysis result for restricted terms detection"""
+    contains_restricted_terms: bool = Field(
+        description="Whether any restricted terms were detected"
+    )
+    detected_matches: List[TermMatch] = Field(
+        default_factory=list,
+        description="List of detected term matches with their variations"
+    )
+    explanation: str = Field(
+        description="Detailed explanation of the analysis"
+    )
+    anonymized_text: Optional[str] = Field(
+        default=None,
+        description="Text with restricted terms replaced with category tags"
+    )
+
+    @property
+    def safe(self) -> bool:
+        return not self.contains_restricted_terms
+
+
+class RestrictedTermsRecognitionResponse(BaseModel):
+    contains_entities: bool
+    detected_entities: Dict[str, List[str]]
+    explanation: str
+    anonymized_text: Optional[str] = None
+
+    @property
+    def safe(self) -> bool:
+        return not self.contains_entities
+
+
+class RestrictedTermsJudge(Guardrail):
+    llm_model: OpenAIModel = Field(default_factory=lambda: OpenAIModel())
+    should_anonymize: bool = False
+
+    def __init__(self, should_anonymize: bool = False, **kwargs):
+        super().__init__(**kwargs)
+        self.should_anonymize = should_anonymize
+        # Pass the OpenAI client to instructor.patch()
+        instructor.patch(client=self.llm_model._openai_client)
+
+    def format_prompts(self, text: str, custom_terms: List[str]) -> tuple[str, str]:
+        user_prompt = f"""Analyze the following text for restricted terms and variations:
+
+<text_to_analyze>
+{text}
+</text_to_analyze>
+
+Look for these specific terms and their variations:
+{', '.join(custom_terms)}
+
+Analyze the text carefully for:
+1. Exact matches
+2. Common misspellings
+3. Abbreviations
+4. Intentional variations (l33t speak, special characters, etc.)
+5. Case variations
+"""
+
+        system_prompt = """You are an expert system for detecting restricted terms, brand mentions, and inappropriate content.
+
+Your task is to:
+1. Identify exact matches of restricted terms
+2. Detect variations including:
+   - Misspellings (both accidental and intentional)
+   - Abbreviations and acronyms
+   - Case variations
+   - L33t speak or special character substitutions
+   - Partial matches within larger words
+
+For each match, you must:
+1. Identify the original restricted term
+2. Note the actual text that matched
+3. Classify the match type
+4. Provide a confidence score
+5. Explain why it's considered a match
+
+Be thorough but avoid false positives. Focus on meaningful matches that indicate actual attempts to use restricted terms.
+
+Return your analysis in the structured format specified by the RestrictedTermsAnalysis model."""
+
+        return user_prompt, system_prompt
+
+    @weave.op()
+    def predict(self, text: str, custom_terms: List[str], **kwargs) -> RestrictedTermsAnalysis:
+        user_prompt, system_prompt = self.format_prompts(text, custom_terms)
+        
+        response = self.llm_model.predict(
+            user_prompts=user_prompt,
+            system_prompt=system_prompt,
+            response_format=RestrictedTermsAnalysis,
+            temperature=0.1,  # Lower temperature for more consistent analysis
+            **kwargs
+        )
+        
+        return response.choices[0].message.parsed
+
+    #TODO: Remove default custom_terms
+    @weave.op()
+    def guard(self, text: str, custom_terms: List[str] = ["Microsoft", "Amazon Web Services", "Facebook", "Meta", "Google", "Salesforce", "Oracle"], aggregate_redaction: bool = True, **kwargs) -> RestrictedTermsRecognitionResponse:
+        """
+        Guard against restricted terms and their variations.
+        
+        Args:
+            text: Text to analyze
+            custom_terms: List of restricted terms to check for
+            
+        Returns:
+            RestrictedTermsRecognitionResponse containing safety assessment and detailed analysis
+        """
+        analysis = self.predict(text, custom_terms, **kwargs)
+        
+        # Create a summary of findings
+        if analysis.contains_restricted_terms:
+            summary_parts = ["Restricted terms detected:"]
+            for match in analysis.detected_matches:
+                summary_parts.append(f"\n- {match.original_term}: {match.matched_text} ({match.match_type})")
+            summary = "\n".join(summary_parts)
+        else:
+            summary = "No restricted terms detected."
+
+        # Updated anonymization logic
+        anonymized_text = None
+        if self.should_anonymize and analysis.contains_restricted_terms:
+            anonymized_text = text
+            for match in analysis.detected_matches:
+                replacement = "[redacted]" if aggregate_redaction else f"[{match.match_type.upper()}]"
+                anonymized_text = anonymized_text.replace(match.matched_text, replacement)
+
+        # Convert detected_matches to a dictionary format
+        detected_entities = {}
+        for match in analysis.detected_matches:
+            if match.original_term not in detected_entities:
+                detected_entities[match.original_term] = []
+            detected_entities[match.original_term].append(match.matched_text)
+
+        return RestrictedTermsRecognitionResponse(
+            contains_entities=analysis.contains_restricted_terms,
+            detected_entities=detected_entities,
+            explanation=summary,
+            anonymized_text=anonymized_text
+        )

guardrails_genie/guardrails/entity_recognition/regex_entity_recognition_guardrail.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel
 
 from ...regex_model import RegexModel
 from ..base import Guardrail
+import re
 
 
 class RegexEntityRecognitionResponse(BaseModel):
@@ -63,19 +64,37 @@ class RegexEntityRecognitionGuardrail(Guardrail):
             should_anonymize=should_anonymize
         )
 
+    def text_to_pattern(self, text: str) -> str:
+        """
+        Convert input text into a regex pattern that matches the exact text.
+        """
+        # Escape special regex characters in the text
+        escaped_text = re.escape(text)
+        # Create a pattern that matches the exact text, case-insensitive
+        return rf"\b{escaped_text}\b"
+
     @weave.op()
-    def guard(self, prompt: str, return_detected_types: bool = True, **kwargs) -> RegexEntityRecognitionResponse | RegexEntityRecognitionSimpleResponse:
+    def guard(self, prompt: str, custom_terms: Optional[list[str]] = None, return_detected_types: bool = True, aggregate_redaction: bool = True, **kwargs) -> RegexEntityRecognitionResponse | RegexEntityRecognitionSimpleResponse:
         """
         Check if the input prompt contains any entities based on the regex patterns.
         
         Args:
             prompt: Input text to check for entities
+            custom_terms: List of custom terms to be converted into regex patterns. If provided, 
+                        only these terms will be checked, ignoring default patterns.
             return_detected_types: If True, returns detailed entity type information
             
         Returns:
             RegexEntityRecognitionResponse or RegexEntityRecognitionSimpleResponse containing detection results
         """
-        result = self.regex_model.check(prompt)
+        if custom_terms:
+            # Create a temporary RegexModel with only the custom patterns
+            temp_patterns = {term: self.text_to_pattern(term) for term in custom_terms}
+            temp_model = RegexModel(patterns=temp_patterns)
+            result = temp_model.check(prompt)
+        else:
+            # Use the original regex_model if no custom terms provided
+            result = self.regex_model.check(prompt)
         
         # Create detailed explanation
         explanation_parts = []
@@ -91,13 +110,13 @@ class RegexEntityRecognitionGuardrail(Guardrail):
             for pattern in result.failed_patterns:
                 explanation_parts.append(f"- {pattern}")
                 
-        # Add anonymization logic
+        # Updated anonymization logic
         anonymized_text = None
         if getattr(self, 'should_anonymize', False) and result.matched_patterns:
             anonymized_text = prompt
             for entity_type, matches in result.matched_patterns.items():
                 for match in matches:
-                    replacement = f"[{entity_type.upper()}]"
+                    replacement = "[redacted]" if aggregate_redaction else f"[{entity_type.upper()}]"
                     anonymized_text = anonymized_text.replace(match, replacement)
         
         if return_detected_types:
@@ -115,5 +134,5 @@ class RegexEntityRecognitionGuardrail(Guardrail):
             )
 
     @weave.op()
-    def predict(self, prompt: str, return_detected_types: bool = True, **kwargs) -> RegexEntityRecognitionResponse | RegexEntityRecognitionSimpleResponse:
-        return self.guard(prompt, return_detected_types=return_detected_types, **kwargs)
+    def predict(self, prompt: str, return_detected_types: bool = True, aggregate_redaction: bool = True, **kwargs) -> RegexEntityRecognitionResponse | RegexEntityRecognitionSimpleResponse:
+        return self.guard(prompt, return_detected_types=return_detected_types, aggregate_redaction=aggregate_redaction, **kwargs)