add: docs for entity recognition guardrails

docs/guardrails/entity_recognition/entity_recognition_guardrails.md

@@ -0,0 +1,136 @@
+# Entity Recognition Guardrails
+
+A collection of guardrails for detecting and anonymizing various types of entities in text, including PII (Personally Identifiable Information), restricted terms, and custom entities.
+
+## Available Guardrails
+
+### 1. Regex Entity Recognition
+Simple pattern-based entity detection using regular expressions.
+
+```python
+from guardrails_genie.guardrails.entity_recognition import RegexEntityRecognitionGuardrail
+
+# Initialize with default PII patterns
+guardrail = RegexEntityRecognitionGuardrail(should_anonymize=True)
+
+# Or with custom patterns
+custom_patterns = {
+    "employee_id": r"EMP\d{6}",
+    "project_code": r"PRJ-[A-Z]{2}-\d{4}"
+}
+guardrail = RegexEntityRecognitionGuardrail(patterns=custom_patterns, should_anonymize=True)
+```
+
+### 2. Presidio Entity Recognition
+Advanced entity detection using Microsoft's Presidio analyzer.
+
+```python
+from guardrails_genie.guardrails.entity_recognition import PresidioEntityRecognitionGuardrail
+
+# Initialize with default entities
+guardrail = PresidioEntityRecognitionGuardrail(should_anonymize=True)
+
+# Or with specific entities
+selected_entities = ["CREDIT_CARD", "US_SSN", "EMAIL_ADDRESS"]
+guardrail = PresidioEntityRecognitionGuardrail(
+    selected_entities=selected_entities,
+    should_anonymize=True
+)
+```
+
+### 3. Transformers Entity Recognition
+Entity detection using transformer-based models.
+
+```python
+from guardrails_genie.guardrails.entity_recognition import TransformersEntityRecognitionGuardrail
+
+# Initialize with default model
+guardrail = TransformersEntityRecognitionGuardrail(should_anonymize=True)
+
+# Or with specific model and entities
+guardrail = TransformersEntityRecognitionGuardrail(
+    model_name="iiiorg/piiranha-v1-detect-personal-information",
+    selected_entities=["GIVENNAME", "SURNAME", "EMAIL"],
+    should_anonymize=True
+)
+```
+
+### 4. LLM Judge for Restricted Terms
+Advanced detection of restricted terms, competitor mentions, and brand protection using LLMs.
+
+```python
+from guardrails_genie.guardrails.entity_recognition import RestrictedTermsJudge
+
+# Initialize with OpenAI model
+guardrail = RestrictedTermsJudge(should_anonymize=True)
+
+# Check for specific terms
+result = guardrail.guard(
+    text="Let's implement features like Salesforce",
+    custom_terms=["Salesforce", "Oracle", "AWS"]
+)
+```
+
+## Usage
+
+All guardrails follow a consistent interface:
+
+```python
+# Initialize a guardrail
+guardrail = RegexEntityRecognitionGuardrail(should_anonymize=True)
+
+# Check text for entities
+result = guardrail.guard("Hello, my email is [email protected]")
+
+# Access results
+print(f"Contains entities: {result.contains_entities}")
+print(f"Detected entities: {result.detected_entities}")
+print(f"Explanation: {result.explanation}")
+print(f"Anonymized text: {result.anonymized_text}")
+```
+
+## Evaluation Tools
+
+The module includes comprehensive evaluation tools and test cases:
+
+- `pii_examples/`: Test cases for PII detection
+- `banned_terms_examples/`: Test cases for restricted terms
+- Benchmark scripts for evaluating model performance
+
+### Running Evaluations
+
+```python
+# PII Detection Benchmark
+from guardrails_genie.guardrails.entity_recognition.pii_examples.pii_benchmark import main
+main()
+
+# (TODO): Restricted Terms Testing
+from guardrails_genie.guardrails.entity_recognition.banned_terms_examples.banned_term_benchmark import main
+main()
+```
+
+## Features
+
+- Entity detection and anonymization
+- Support for multiple detection methods (regex, Presidio, transformers, LLMs)
+- Customizable entity types and patterns
+- Detailed explanations of detected entities
+- Comprehensive evaluation framework
+- Support for custom terms and patterns
+- Batch processing capabilities
+- Performance metrics and benchmarking
+
+## Response Format
+
+All guardrails return responses with the following structure:
+
+```python
+{
+    "contains_entities": bool,
+    "detected_entities": {
+        "entity_type": ["detected_value_1", "detected_value_2"]
+    },
+    "explanation": str,
+    "anonymized_text": Optional[str]
+}
+```

docs/guardrails/entity_recognition/llm_judge_entity_recognition_guardrail.md

@@ -0,0 +1,3 @@
+# LLM Judge for Entity Recognition Guardrail
+
+::: guardrails_genie.guardrails.entity_recognition.llm_judge_entity_recognition_guardrail

docs/guardrails/entity_recognition/presidio_entity_recognition_guardrail.md

@@ -0,0 +1,3 @@
+# Presidio Entity Recognition Guardrail
+
+::: guardrails_genie.guardrails.entity_recognition.presidio_entity_recognition_guardrail

docs/guardrails/entity_recognition/regex_entity_recognition_guardrail.md

@@ -0,0 +1,3 @@
+# Regex Entity Recognition Guardrail
+
+::: guardrails_genie.guardrails.entity_recognition.regex_entity_recognition_guardrail

docs/guardrails/entity_recognition/transformers_entity_recognition_guardrail.md

@@ -0,0 +1,3 @@
+# Transformers Entity Recognition Guardrail
+
+::: guardrails_genie.guardrails.entity_recognition.transformers_entity_recognition_guardrail

guardrails_genie/guardrails/entity_recognition/llm_judge_entity_recognition_guardrail.py

@@ -54,6 +54,36 @@ class RestrictedTermsRecognitionResponse(BaseModel):
 
 
 class RestrictedTermsJudge(Guardrail):
+    """
+    A class to detect and analyze restricted terms and their variations in text using an LLM model.
+
+    The RestrictedTermsJudge class extends the Guardrail class and utilizes an OpenAIModel
+    to identify restricted terms and their variations within a given text. It provides
+    functionality to format prompts for the LLM, predict restricted terms, and optionally
+    anonymize detected terms in the text.
+
+    !!! example "Using RestrictedTermsJudge"
+        ```python
+        from guardrails_genie.guardrails.entity_recognition import RestrictedTermsJudge
+
+        # Initialize with OpenAI model
+        guardrail = RestrictedTermsJudge(should_anonymize=True)
+
+        # Check for specific terms
+        result = guardrail.guard(
+            text="Let's implement features like Salesforce",
+            custom_terms=["Salesforce", "Oracle", "AWS"]
+        )
+        ```
+
+    Attributes:
+        llm_model (OpenAIModel): An instance of OpenAIModel used for predictions.
+        should_anonymize (bool): A flag indicating whether detected terms should be anonymized.
+
+    Args:
+        should_anonymize (bool): A flag indicating whether detected terms should be anonymized.
+    """
+
     llm_model: OpenAIModel = Field(default_factory=lambda: OpenAIModel())
     should_anonymize: bool = False
 
@@ -139,14 +169,32 @@ Return your analysis in the structured format specified by the RestrictedTermsAn
         **kwargs,
     ) -> RestrictedTermsRecognitionResponse:
         """
-        Guard against restricted terms and their variations.
+        Analyzes the provided text to identify and handle restricted terms and their variations.
+
+        This function utilizes a predictive model to scan the input text for any occurrences of
+        specified restricted terms, including their variations such as misspellings, abbreviations,
+        and case differences. It returns a detailed analysis of the findings, including whether
+        restricted terms were detected, a summary of the matches, and an optional anonymized version
+        of the text.
+
+        The function operates by first calling the `predict` method to perform the analysis based on
+        the given text and custom terms. If restricted terms are found, it constructs a summary of
+        these findings. Additionally, if anonymization is enabled, it replaces detected terms in the
+        text with a redacted placeholder or a specific match type indicator, depending on the
+        `aggregate_redaction` flag.
 
         Args:
-            text: Text to analyze
-            custom_terms: List of restricted terms to check for
+            text (str): The text to be analyzed for restricted terms.
+            custom_terms (List[str]): A list of restricted terms to check against the text. Defaults
+                to a predefined list of company names.
+            aggregate_redaction (bool): Determines the anonymization strategy. If True, all matches
+                are replaced with "[redacted]". If False, matches are replaced
+                with their match type in uppercase.
 
         Returns:
-            RestrictedTermsRecognitionResponse containing safety assessment and detailed analysis
+            RestrictedTermsRecognitionResponse: An object containing the results of the analysis,
+                including whether restricted terms were found, a dictionary of detected entities,
+                a summary explanation, and the anonymized text if applicable.
         """
         analysis = self.predict(text, custom_terms, **kwargs)
 

guardrails_genie/guardrails/entity_recognition/presidio_entity_recognition_guardrail.py

@@ -36,6 +36,49 @@ class PresidioEntityRecognitionSimpleResponse(BaseModel):
 
 # TODO: Add support for transformers workflow and not just Spacy
 class PresidioEntityRecognitionGuardrail(Guardrail):
+    """
+    A guardrail class for entity recognition and anonymization using Presidio.
+
+    This class extends the Guardrail base class to provide functionality for
+    detecting and optionally anonymizing entities in text using the Presidio
+    library. It leverages Presidio's AnalyzerEngine and AnonymizerEngine to
+    perform these tasks.
+
+    !!! example "Using PresidioEntityRecognitionGuardrail"
+        ```python
+        from guardrails_genie.guardrails.entity_recognition import PresidioEntityRecognitionGuardrail
+
+        # Initialize with default entities
+        guardrail = PresidioEntityRecognitionGuardrail(should_anonymize=True)
+
+        # Or with specific entities
+        selected_entities = ["CREDIT_CARD", "US_SSN", "EMAIL_ADDRESS"]
+        guardrail = PresidioEntityRecognitionGuardrail(
+            selected_entities=selected_entities,
+            should_anonymize=True
+        )
+        ```
+
+    Attributes:
+        analyzer (AnalyzerEngine): The Presidio engine used for entity analysis.
+        anonymizer (AnonymizerEngine): The Presidio engine used for text anonymization.
+        selected_entities (List[str]): A list of entity types to detect in the text.
+        should_anonymize (bool): A flag indicating whether detected entities should be anonymized.
+        language (str): The language of the text to be analyzed.
+
+    Args:
+        selected_entities (Optional[List[str]]): A list of entity types to detect in the text.
+        should_anonymize (bool): A flag indicating whether detected entities should be anonymized.
+        language (str): The language of the text to be analyzed.
+        deny_lists (Optional[Dict[str, List[str]]]): A dictionary of entity types and their
+            corresponding deny lists.
+        regex_patterns (Optional[Dict[str, List[Dict[str, str]]]]): A dictionary of entity
+            types and their corresponding regex patterns.
+        custom_recognizers (Optional[List[Any]]): A list of custom recognizers to add to the
+            analyzer.
+        show_available_entities (bool): A flag indicating whether to print available entities.
+    """
+
     @staticmethod
     def get_available_entities() -> List[str]:
         registry = RecognizerRegistry()
@@ -137,11 +180,24 @@ class PresidioEntityRecognitionGuardrail(Guardrail):
         self, prompt: str, return_detected_types: bool = True, **kwargs
     ) -> PresidioEntityRecognitionResponse | PresidioEntityRecognitionSimpleResponse:
         """
-        Check if the input prompt contains any entities using Presidio.
+        Analyzes the input prompt for entity recognition using the Presidio framework.
+
+        This function utilizes the Presidio AnalyzerEngine to detect entities within the
+        provided text prompt. It supports custom recognizers, deny lists, and regex patterns
+        for entity detection. The detected entities are grouped by their types and an
+        explanation of the findings is generated. If anonymization is enabled, the detected
+        entities in the text are anonymized.
 
         Args:
-            prompt: The text to analyze
-            return_detected_types: If True, returns detailed entity type information
+            prompt (str): The text to be analyzed for entity recognition.
+            return_detected_types (bool): Determines the type of response. If True, the
+                response includes detailed information about detected entity types.
+
+        Returns:
+            PresidioEntityRecognitionResponse | PresidioEntityRecognitionSimpleResponse:
+            A response object containing information about whether entities were detected,
+            the types and instances of detected entities, an explanation of the analysis,
+            and optionally, the anonymized text if anonymization is enabled.
         """
         # Analyze text for entities
         analyzer_results = self.analyzer.analyze(

guardrails_genie/guardrails/entity_recognition/regex_entity_recognition_guardrail.py

@@ -30,6 +30,40 @@ class RegexEntityRecognitionSimpleResponse(BaseModel):
 
 
 class RegexEntityRecognitionGuardrail(Guardrail):
+    """
+    A guardrail class for recognizing and optionally anonymizing entities in text using regular expressions.
+
+    This class extends the Guardrail base class and utilizes a RegexModel to detect entities in the input text
+    based on predefined or custom regex patterns. It provides functionality to check for entities, anonymize
+    detected entities, and return detailed information about the detected entities.
+
+    !!! example "Using RegexEntityRecognitionGuardrail"
+        ```python
+        from guardrails_genie.guardrails.entity_recognition import RegexEntityRecognitionGuardrail
+
+        # Initialize with default PII patterns
+        guardrail = RegexEntityRecognitionGuardrail(should_anonymize=True)
+
+        # Or with custom patterns
+        custom_patterns = {
+            "employee_id": r"EMP\d{6}",
+            "project_code": r"PRJ-[A-Z]{2}-\d{4}"
+        }
+        guardrail = RegexEntityRecognitionGuardrail(patterns=custom_patterns, should_anonymize=True)
+        ```
+
+    Attributes:
+        regex_model (RegexModel): An instance of RegexModel used for entity recognition.
+        patterns (Dict[str, str]): A dictionary of regex patterns for entity recognition.
+        should_anonymize (bool): A flag indicating whether detected entities should be anonymized.
+        DEFAULT_PATTERNS (ClassVar[Dict[str, str]]): A dictionary of default regex patterns for common entities.
+
+    Args:
+        use_defaults (bool): If True, use default patterns. If False, use custom patterns.
+        should_anonymize (bool): If True, anonymize detected entities.
+        show_available_entities (bool): If True, print available entity types.
+    """
+
     regex_model: RegexModel
     patterns: Dict[str, str] = {}
     should_anonymize: bool = False
@@ -107,16 +141,29 @@ class RegexEntityRecognitionGuardrail(Guardrail):
         **kwargs,
     ) -> RegexEntityRecognitionResponse | RegexEntityRecognitionSimpleResponse:
         """
-        Check if the input prompt contains any entities based on the regex patterns.
+        Analyzes the input prompt to detect entities based on predefined or custom regex patterns.
+
+        This function checks the provided text (prompt) for entities using regex patterns. It can
+        utilize either default patterns or custom terms provided by the user. If custom terms are
+        specified, they are converted into regex patterns, and only these are used for entity detection.
+        The function returns detailed information about detected entities and can optionally anonymize
+        the detected entities in the text.
 
         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
+            prompt (str): The input text to be analyzed for entity detection.
+            custom_terms (Optional[list[str]]): A list of custom terms to be converted into regex patterns.
+                If provided, only these terms will be checked, ignoring default patterns.
+            return_detected_types (bool): If True, the function returns detailed information about the
+                types of entities detected in the text.
+            aggregate_redaction (bool): Determines the anonymization strategy. If True, all detected
+                entities are replaced with a generic "[redacted]" label. If False, each entity type is
+                replaced with its specific label (e.g., "[ENTITY_TYPE]").
 
         Returns:
-            RegexEntityRecognitionResponse or RegexEntityRecognitionSimpleResponse containing detection results
+            RegexEntityRecognitionResponse or RegexEntityRecognitionSimpleResponse: An object containing
+            the results of the entity detection, including whether entities were found, the types and
+            counts of detected entities, an explanation of the detection process, and optionally, the
+            anonymized text.
         """
         if custom_terms:
             # Create a temporary RegexModel with only the custom patterns

guardrails_genie/guardrails/entity_recognition/transformers_entity_recognition_guardrail.py

@@ -29,7 +29,40 @@ class TransformersEntityRecognitionSimpleResponse(BaseModel):
 
 
 class TransformersEntityRecognitionGuardrail(Guardrail):
-    """Generic guardrail for detecting entities using any token classification model."""
+    """Generic guardrail for detecting entities using any token classification model.
+
+    This class leverages a transformer-based token classification model to detect and
+    optionally anonymize entities in a given text. It uses the HuggingFace `transformers`
+    library to load a pre-trained model and perform entity recognition.
+
+    !!! example "Using TransformersEntityRecognitionGuardrail"
+        ```python
+        from guardrails_genie.guardrails.entity_recognition import TransformersEntityRecognitionGuardrail
+
+        # Initialize with default model
+        guardrail = TransformersEntityRecognitionGuardrail(should_anonymize=True)
+
+        # Or with specific model and entities
+        guardrail = TransformersEntityRecognitionGuardrail(
+            model_name="iiiorg/piiranha-v1-detect-personal-information",
+            selected_entities=["GIVENNAME", "SURNAME", "EMAIL"],
+            should_anonymize=True
+        )
+        ```
+
+    Attributes:
+        _pipeline (Optional[object]): The transformer pipeline for token classification.
+        selected_entities (List[str]): List of entities to detect.
+        should_anonymize (bool): Flag indicating whether detected entities should be anonymized.
+        available_entities (List[str]): List of all available entities that the model can detect.
+
+    Args:
+        model_name (str): The name of the pre-trained model to use for entity recognition.
+        selected_entities (Optional[List[str]]): A list of specific entities to detect.
+            If None, all available entities will be used.
+        should_anonymize (bool): If True, detected entities will be anonymized.
+        show_available_entities (bool): If True, available entity types will be printed.
+    """
 
     _pipeline: Optional[object] = None
     selected_entities: List[str]
@@ -161,12 +194,26 @@ class TransformersEntityRecognitionGuardrail(Guardrail):
         TransformersEntityRecognitionResponse
         | TransformersEntityRecognitionSimpleResponse
     ):
-        """Check if the input prompt contains any entities using the transformer pipeline.
+        """Analyze the input prompt for entity recognition and optionally anonymize detected entities.
+
+        This function utilizes a transformer-based pipeline to detect entities within the provided
+        text prompt. It returns a response indicating whether any entities were found, along with
+        detailed information about the detected entities if requested. The function can also anonymize
+        the detected entities in the text based on the specified parameters.
 
         Args:
-            prompt: The text to analyze
-            return_detected_types: If True, returns detailed entity type information
-            aggregate_redaction: If True, uses generic [redacted] instead of entity type
+            prompt (str): The text to be analyzed for entity detection.
+            return_detected_types (bool): If True, the response includes detailed information about
+                the types of entities detected. Defaults to True.
+            aggregate_redaction (bool): If True, detected entities are anonymized using a generic
+                [redacted] marker. If False, the specific entity type is used in the redaction.
+                Defaults to True.
+
+        Returns:
+            TransformersEntityRecognitionResponse or TransformersEntityRecognitionSimpleResponse:
+            A response object containing information about the presence of entities, an explanation
+            of the detection process, and optionally, the anonymized text if entities were detected
+            and anonymization is enabled.
         """
         # Detect entities
         detected_entities = self._detect_entities(prompt)

mkdocs.yml

@@ -62,6 +62,12 @@ nav:
   - Guardrails:
     - Guardrail Base Class: 'guardrails/base.md'
     - Guardrail Manager: 'guardrails/manager.md'
+    - Entity Recognition Guardrails:
+      - About: 'guardrails/entity_recognition/entity_recognition_guardrails.md'
+      - Regex Entity Recognition Guardrail: 'guardrails/entity_recognition/regex_entity_recognition_guardrail.md'
+      - Presidio Entity Recognition Guardrail: 'guardrails/entity_recognition/presidio_entity_recognition_guardrail.md'
+      - Transformers Entity Recognition Guardrail: 'guardrails/entity_recognition/transformers_entity_recognition_guardrail.md'
+      - LLM Judge for Entity Recognition Guardrail: 'guardrails/entity_recognition/llm_judge_entity_recognition_guardrail.md'
     - Prompt Injection Guardrails:
       - Classifier Guardrail: 'guardrails/prompt_injection/classifier.md'
       - Survey Guardrail: 'guardrails/prompt_injection/llm_survey.md'

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "torch>=2.5.1",
     "presidio-analyzer>=2.2.355",
     "presidio-anonymizer>=2.2.355",
+    "instructor>=1.7.0",
 ]
 
 [project.optional-dependencies]