Add Private AI Integration

docs/user_guides/community/privateai.md

@@ -0,0 +1,61 @@
+# 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.
+
+## 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:
+
+```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:
+      - detect pii on input
+  output:
+    flows:
+      - detect pii on output
+```
+
+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:
+
+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.
+
+The `detect_pii` action in `nemoguardrails/library/privateai/actions.py` handles the PII detection process.
+
+## Customization
+
+You can customize the PII detection behavior by modifying the `entities` lists in the `config.yml` file. Refer to the Private AI documentation for a complete list of [supported entity types](https://docs.private-ai.com/entities/?utm_medium=github&utm_campaign=nemo-guardrails).
+
+## Error Handling
+
+If the Private AI detection API request fails, the system will assume PII is present as a precautionary measure.
+
+## 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.
+
+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

@@ -22,6 +22,7 @@ NeMo Guardrails comes with a library of built-in guardrails that you can easily
    - [AutoAlign](#autoalign)
    - [Cleanlab Trustworthiness Score](#cleanlab)
    - [GCP Text Moderation](#gcp-text-moderation)
+   - [Private AI PII detection](#private-ai-pii-detection)
    - OpenAI Moderation API - *[COMING SOON]*
 
 4. Other
@@ -670,6 +671,46 @@ rails:
 
 For more details, check out the [GCP Text Moderation](./community/gcp-text-moderations.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
+
+```yaml
+rails:
+  input:
+    flows:
+      - detect pii on input
+  output:
+    flows:
+      - detect pii on output
+  retrieval:
+    flows:
+      - detect pii on retrieval
+```
+
+For more details, check out the [Private AI Integration](./community/privateai.md) page.
+
 ## Other
 
 ### Jailbreak Detection Heuristics

examples/configs/privateai/README.md

@@ -0,0 +1,11 @@
+# Private AI Configuration Example
+
+This example contains configuration files for using Private AI in your NeMo Guardrails project.
+
+For more details on the Private AI integration, see [Private AI Integration User Guide](../../../docs/user_guides/community/privateai.md).
+
+## Structure
+
+The Private AI configuration example is organized as follows:
+
+1. [pii_detection](./pii_detection) - Configuration for using Private AI for PII detection.

examples/configs/privateai/pii_detection/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 detected by default.
+          - NAME_FAMILY
+          - LOCATION_ADDRESS_STREET
+          - EMAIL_ADDRESS
+  input:
+    flows:
+      - detect pii on input
+
+  output:
+    flows:
+      - detect pii on output

examples/notebooks/privateai_pii_detection.ipynb

@@ -0,0 +1,186 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Private AI PII detection example\n",
+    "\n",
+    "This notebook shows how to use Private AI for PII detection in NeMo Guardrails."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from nemoguardrails import LLMRails, RailsConfig"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create rails with Private AI PII detection\n",
+    "\n",
+    "For this step you'll need your OpenAI API key & Private AI API key.\n",
+    "\n",
+    "You can get your Private AI API key by signing up on the [Private AI Portal](https://portal.private-ai.com). For more details on Private AI integration, check out this [user guide](../../docs/user_guides/community/privateai.md).\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"PAI_API_KEY\"] = \"YOUR PRIVATE AI API KEY\"  # Visit https://portal.private-ai.com to get your API key\n",
+    "\n",
+    "YAML_CONFIG = \"\"\"\n",
+    "models:\n",
+    "  - type: main\n",
+    "    engine: openai\n",
+    "    model: gpt-3.5-turbo-instruct\n",
+    "\n",
+    "rails:\n",
+    "  config:\n",
+    "    privateai:\n",
+    "      server_endpoint: https://api.private-ai.com/cloud/v3/process/text\n",
+    "      input:\n",
+    "        entities:\n",
+    "          - NAME_FAMILY\n",
+    "          - LOCATION_ADDRESS_STREET\n",
+    "          - EMAIL_ADDRESS\n",
+    "      output:\n",
+    "        entities:\n",
+    "          - NAME_FAMILY\n",
+    "          - LOCATION_ADDRESS_STREET\n",
+    "          - EMAIL_ADDRESS\n",
+    "  input:\n",
+    "    flows:\n",
+    "      - detect pii on input\n",
+    "\n",
+    "  output:\n",
+    "    flows:\n",
+    "      - detect pii on output\n",
+    "\"\"\"\n",
+    "\n",
+    "\n",
+    "\n",
+    "config = RailsConfig.from_content(yaml_content=YAML_CONFIG)\n",
+    "rails = LLMRails(config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Input rails"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"Hello! I'm John. My email id is text@gmail.com. I live in California, USA.\"}])\n",
+    "\n",
+    "info = rails.explain()\n",
+    "\n",
+    "print(\"Response\")\n",
+    "print(\"----------------------------------------\")\n",
+    "print(response[\"content\"])\n",
+    "\n",
+    "\n",
+    "print(\"\\n\\nColang history\")\n",
+    "print(\"----------------------------------------\")\n",
+    "print(info.colang_history)\n",
+    "\n",
+    "print(\"\\n\\nLLM calls summary\")\n",
+    "print(\"----------------------------------------\")\n",
+    "info.print_llm_calls_summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Output rails"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"give me a sample email id\"}])\n",
+    "\n",
+    "info = rails.explain()\n",
+    "\n",
+    "print(\"Response\")\n",
+    "print(\"----------------------------------------\\n\\n\")\n",
+    "print(response[\"content\"])\n",
+    "\n",
+    "\n",
+    "print(\"\\n\\nColang history\")\n",
+    "print(\"----------------------------------------\")\n",
+    "print(info.colang_history)\n",
+    "\n",
+    "print(\"\\n\\nLLM calls summary\")\n",
+    "print(\"----------------------------------------\")\n",
+    "info.print_llm_calls_summary()\n",
+    "\n",
+    "\n",
+    "print(\"\\n\\nCompletions where PII was detected!\")\n",
+    "print(\"----------------------------------------\")\n",
+    "print(info.llm_calls[0].completion)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "nemo",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

nemoguardrails/library/privateai/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.