feat: add pii masking capability to private ai integration (#901)

* Add tests for private ai pii masking

* Add pii masking action and flow for private ai

* Update private ai docs to add pii masking

* Add private ai pii masking example config

* Add private ai pii masking example in notebook

* Apply suggestions from code review

Co-authored-by: Pouyan <13303554+Pouyanpi@users.noreply.github.com>
Signed-off-by: Girish Sharma <girishsharma001@gmail.com>

* Test private ai with live api instead of mocked detector

* Improve error handling for private ai

---------

Signed-off-by: Girish Sharma <girishsharma001@gmail.com>
Co-authored-by: Pouyan <13303554+Pouyanpi@users.noreply.github.com>

Author: letmerecall

docs/user-guides/community/privateai.md

@@ -1,13 +1,15 @@
 # Private AI Integration
 
-[Private AI](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails) allows you to detect and mask Personally Identifiable Information (PII) in your data. This integration enables NeMo Guardrails to use Private AI for PII detection in input, output and retrieval flows.
+[Private AI](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails) allows you to detect and mask Personally Identifiable Information (PII) in your data. This integration enables NeMo Guardrails to use Private AI for PII detection and masking in input, output, and retrieval flows.
 
 ## Setup
 
 1. Ensure that you have access to Private AI API server running locally or in the cloud. To get started with the cloud version, you can use the [Private AI Portal](https://portal.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails). For containerized deployments, check out this [Quickstart Guide](https://docs.private-ai.com/quickstart/?utm_medium=github&utm_campaign=nemo-guardrails).
 
 2. Update your `config.yml` file to include the Private AI settings:
 
+**PII detection config**
+
 ```yaml
 rails:
   config:
@@ -31,19 +33,48 @@ rails:
       - detect pii on output
 ```
 
+The detection flow will not let the input/output/retrieval text pass if PII is detected.
+
+**PII masking config**
+
+```yaml
+rails:
+  config:
+    privateai:
+      server_endpoint: http://your-privateai-api-endpoint/process/text  # Replace this with your Private AI process text endpoint
+      input:
+        entities:  # If no entity is specified here, all supported entities will be detected by default.
+          - NAME_FAMILY
+          - LOCATION_ADDRESS_STREET
+          - EMAIL_ADDRESS
+      output:
+        entities:
+          - NAME_FAMILY
+          - LOCATION_ADDRESS_STREET
+          - EMAIL_ADDRESS
+  input:
+    flows:
+      - mask pii on input
+  output:
+    flows:
+      - mask pii on output
+```
+
+The masking flow will mask the PII in the input/output/retrieval text before they are sent to the LLM/user. For example, `Hi John Doe, my email is john.doe@example.com` will be converted to `Hi [NAME], my email is [EMAIL_ADDRESS]`.
+
 Replace `http://your-privateai-api-endpoint/process/text` with your actual Private AI process text endpoint and set the `PAI_API_KEY` environment variable if you're using the Private AI cloud API.
 
 3. You can customize the `entities` list under both `input` and `output` to include the PII types you want to detect. A full list of supported entities can be found [here](https://docs.private-ai.com/entities/?utm_medium=github&utm_campaign=nemo-guardrails).
 
 ## Usage
 
-Once configured, the Private AI integration will automatically:
+Once configured, the Private AI integration can automatically:
 
-1. Detect PII in user inputs before they are processed by the LLM.
-2. Detect PII in LLM outputs before they are sent back to the user.
-3. Detect PII in retrieved chunks before they are sent to the LLM.
+1. Detect or mask PII in user inputs before they are processed by the LLM.
+2. Detect or mask PII in LLM outputs before they are sent back to the user.
+3. Detect or mask PII in retrieved chunks before they are sent to the LLM.
 
-The `detect_pii` action in `nemoguardrails/library/privateai/actions.py` handles the PII detection process.
+The `detect_pii` and `mask_pii` actions in `nemoguardrails/library/privateai/actions.py` handle the PII detection and masking processes, respectively.
 
 ## Customization
 
@@ -56,6 +87,6 @@ If the Private AI detection API request fails, the system will assume PII is pre
 ## Notes
 
 - Ensure that your Private AI process text endpoint is properly set up and accessible from your NeMo Guardrails environment.
-- The integration currently supports PII detection only.
+- The integration currently supports PII detection and masking.
 
 For more information on Private AI and its capabilities, please refer to the [Private AI documentation](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails).

docs/user-guides/guardrails-library.md

@@ -750,9 +750,9 @@ For more details, check out the [GCP Text Moderation](https://github.com/NVIDIA/
 
 ### Private AI PII Detection
 
-NeMo Guardrails supports using [Private AI API](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails) for PII detection in input, output and retrieval flows.
+NeMo Guardrails supports using [Private AI API](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails) for PII detection and masking input, output and retrieval flows.
 
-To activate the PII detection, you need specify `server_endpoint`, and the entities that you want to detect. You'll also need to set the `PAI_API_KEY` environment variable if you're using the Private AI cloud API.
+To activate the PII detection or masking, you need specify `server_endpoint`, and the entities that you want to detect or mask. You'll also need to set the `PAI_API_KEY` environment variable if you're using the Private AI cloud API.
 
 ```yaml
 rails:
@@ -773,6 +773,8 @@ rails:
 
 #### Example usage
 
+**PII detection**
+
 ```yaml
 rails:
   input:
@@ -786,44 +788,19 @@ rails:
       - detect pii on retrieval
 ```
 
-For more details, check out the [Private AI Integration](https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/docs/user-guides/community/privateai.md) page.
-
-### Private AI PII Detection
-
-NeMo Guardrails supports using [Private AI API](https://docs.private-ai.com/?utm_medium=github&utm_campaign=nemo-guardrails) for PII detection in input, output and retrieval flows.
-
-To activate the PII detection, you need specify `server_endpoint`, and the entities that you want to detect. You'll also need to set the `PAI_API_KEY` environment variable if you're using the Private AI cloud API.
-
-```yaml
-rails:
-  config:
-    privateai:
-      server_endpoint: http://your-privateai-api-endpoint/process/text  # Replace this with your Private AI process text endpoint
-      input:
-        entities:  # If no entity is specified here, all supported entities will be detected by default.
-          - NAME_FAMILY
-          - EMAIL_ADDRESS
-          ...
-      output:
-        entities:
-          - NAME_FAMILY
-          - EMAIL_ADDRESS
-          ...
-```
-
-#### Example usage
+**PII masking**
 
 ```yaml
 rails:
   input:
     flows:
-      - detect pii on input
+      - mask pii on input
   output:
     flows:
-      - detect pii on output
+      - mask pii on output
   retrieval:
     flows:
-      - detect pii on retrieval
+      - mask pii on retrieval
 ```
 
 For more details, check out the [Private AI Integration](./community/privateai.md) page.

examples/configs/privateai/pii_masking/config.yml

@@ -0,0 +1,26 @@
+models:
+  - type: main
+    engine: openai
+    model: gpt-3.5-turbo-instruct
+
+rails:
+  config:
+    privateai:
+      server_endpoint: https://api.private-ai.com/cloud/v3/process/text
+      input:
+        entities:
+          - NAME_FAMILY
+          - LOCATION_ADDRESS_STREET
+          - EMAIL_ADDRESS
+      output:
+        entities:   # If no entity is specified here, all supported entities will be masked by default.
+          - NAME_FAMILY
+          - LOCATION_ADDRESS_STREET
+          - EMAIL_ADDRESS
+  input:
+    flows:
+      - mask pii on input
+
+  output:
+    flows:
+      - mask pii on output