Adding entity list to PII

.github/workflows/docs.yml

@@ -24,30 +24,31 @@ jobs:
       url: ${{ steps.deployment.outputs.page_url }}
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
 
-      - name: Install Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: 3.9
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@1e862dfacbd1d6d858c55d9b792c756523627244 # v7.1.4
+
+      - name: Set up Python
+        run: uv python install
 
       - name: Install dependencies
-        run: pip install .[dev]
+        run: uv sync --extra dev
 
       - name: Build docs
         run: |
           cd docs
-          make html
+          uv run make html
 
       - name: Setup Pages
-        uses: actions/configure-pages@v5
+        uses: actions/configure-pages@983d7736d9b0ae728b81ab479565c72886d7745b # v5.0.0
 
       - name: Upload artifact
-        uses: actions/upload-pages-artifact@v3
+        uses: actions/upload-pages-artifact@7b1f4a764d45c48632c6b24a0339c27f5614fb0b # v4.0.0
         with:
           path: 'docs/build/html'
 
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
-
+        uses: actions/deploy-pages@d6db90164ac5ed86f2b6aed7e0febac5b3c0c03e # v4.0.5

.github/workflows/main.yml

@@ -11,13 +11,19 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
 
-      - name: Install Pytest
-        run: pip install pytest
+      - name: Install uv
+        uses: astral-sh/setup-uv@1e862dfacbd1d6d858c55d9b792c756523627244 # v7.1.4
+
+      - name: Set up Python
+        run: uv python install
+
+      - name: Install dependencies
+        run: uv sync --extra dev
 
       - name: Run Tests
-        run: pytest tests
+        run: uv run pytest tests
         env:
           PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY }}
           PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL }}
@@ -27,8 +33,11 @@ jobs:
           TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}
           TEST_RERANK_MODEL: ${{ vars.TEST_RERANK_MODEL }}
 
-      - name: To PyPI using Flit
-        uses: AsifArmanRahman/to-pypi-using-flit@v1
-        with:
-          password: ${{ secrets.TESTPYPI_API_TOKEN }}
-          repository-url: https://test.pypi.org/legacy/
+      - name: Build package
+        run: uv build
+
+      - name: Publish to TestPyPI
+        env:
+          UV_PUBLISH_TOKEN: ${{ secrets.TESTPYPI_API_TOKEN }}
+          UV_PUBLISH_URL: https://test.pypi.org/legacy/
+        run: uv publish

.github/workflows/pr.yml

@@ -6,27 +6,32 @@ on:
       - main
 
 jobs:
-
   tests:
     name: Installing requirements and testing functions
     runs-on: ubuntu-latest
     steps:
-    - name: Checkout
-      uses: actions/checkout@v2
+      - name: Checkout
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@1e862dfacbd1d6d858c55d9b792c756523627244 # v7.1.4
+
+      - name: Set up Python
+        run: uv python install
 
-    - name: Install Dependencies
-      run: pip install pytest ruff
+      - name: Install dependencies
+        run: uv sync --extra dev
 
-    - name: Run Ruff
-      run: ruff check --ignore F401 --output-format=github .
+      - name: Run Ruff
+        run: uv run ruff check --ignore F401 --output-format=github .
 
-    - name: Run Tests
-      run: pytest tests
-      env:
-        PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY}}
-        PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL}}
-        TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
-        TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
-        TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
-        TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}
-        TEST_RERANK_MODEL: ${{ vars.TEST_RERANK_MODEL }}
+      - name: Run Tests
+        run: uv run pytest tests
+        env:
+          PREDICTIONGUARD_API_KEY: ${{ secrets.PREDICTIONGUARD_API_KEY }}
+          PREDICTIONGUARD_URL: ${{ vars.PREDICTIONGUARD_URL }}
+          TEST_CHAT_MODEL: ${{ vars.TEST_CHAT_MODEL }}
+          TEST_TEXT_EMBEDDINGS_MODEL: ${{ vars.TEST_TEXT_EMBEDDINGS_MODEL }}
+          TEST_MULTIMODAL_EMBEDDINGS_MODEL: ${{ vars.TEST_MULTIMODAL_EMBEDDINGS_MODEL }}
+          TEST_VISION_MODEL: ${{ vars.TEST_VISION_MODEL }}
+          TEST_RERANK_MODEL: ${{ vars.TEST_RERANK_MODEL }}

.github/workflows/pypi-release.yml

@@ -11,9 +11,18 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
 
-      - name: Publish to PyPI using Flit
-        uses: AsifArmanRahman/to-pypi-using-flit@v1
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
+      - name: Install uv
+        uses: astral-sh/setup-uv@1e862dfacbd1d6d858c55d9b792c756523627244 # v7.1.4
+
+      - name: Set up Python
+        run: uv python install
+
+      - name: Build package
+        run: uv build
+
+      - name: Publish to PyPI
+        env:
+          UV_PUBLISH_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
+        run: uv publish

.gitignore

@@ -107,6 +107,3 @@ venv.bak/
 
 # JetBrains Folder
 .idea
-
-# uv package management
-uv.lock

predictionguard/src/audio.py

@@ -65,6 +65,7 @@ def create(
         toxicity: Optional[bool] = False,
         pii: Optional[str] = "",
         replace_method: Optional[str] = "",
+        entity_list: Optional[List[str]] = None,
         injection: Optional[bool] = False,
     ) -> Dict[str, Any]:
         """
@@ -81,6 +82,7 @@ def create(
         :param toxicity: Whether to check for output toxicity
         :param pii: Whether to check for or replace pii
         :param replace_method: Replace method for any PII that is present.
+        :param entity_list: List of entities to ignore in the PII check.
         :param injection: Whether to check for prompt injection
         :result: A dictionary containing the transcribed text.
         """
@@ -98,6 +100,7 @@ def create(
             response_format,
             pii,
             replace_method,
+            entity_list,
             injection,
             toxicity,
         )
@@ -118,6 +121,7 @@ def _transcribe_audio(
             response_format,
             pii,
             replace_method,
+            entity_list,
             injection,
             toxicity,
     ):
@@ -131,6 +135,7 @@ def _transcribe_audio(
             "Toxicity": str(toxicity),
             "Pii": str(pii),
             "Replace-Method": str(replace_method),
+            "Entity-List": str(entity_list),
             "Injection": str(injection)
         }
 

predictionguard/src/documents.py

@@ -1,5 +1,5 @@
 import requests
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from ..version import __version__
 
@@ -61,6 +61,7 @@ def create(
         toxicity: Optional[bool] = False,
         pii: Optional[str] = "",
         replace_method: Optional[str] = "",
+        entity_list: Optional[List[str]] = None,
         injection: Optional[bool] = False,
     ) -> Dict[str, Any]:
         """
@@ -75,6 +76,7 @@ def create(
         :param toxicity: Whether to check for output toxicity
         :param pii: Whether to check for or replace pii
         :param replace_method: Replace method for any PII that is present.
+        :param entity_list: List of entities to ignore in the PII check.
         :param injection: Whether to check for prompt injection
         :result: A dictionary containing the title, content, and length of the document.
         """
@@ -83,15 +85,15 @@ def create(
         choices = self._extract_documents(
             file, embed_images, output_format,
             chunk_document, chunk_size, enable_ocr,
-            toxicity, pii, replace_method, injection
+            toxicity, pii, replace_method, entity_list, injection
         )
         return choices
 
     def _extract_documents(
             self, file, embed_images,
             output_format, chunk_document,
             chunk_size, enable_ocr, toxicity,
-            pii, replace_method, injection
+            pii, replace_method, entity_list, injection
     ):
         """
         Function to extract a document.
@@ -103,6 +105,7 @@ def _extract_documents(
             "Toxicity": str(toxicity),
             "Pii": str(pii),
             "Replace-Method": str(replace_method),
+            "Entity-List": str(entity_list),
             "Injection": str(injection)
         }
 

predictionguard/src/injection.py

@@ -1,7 +1,7 @@
 import json
 
 import requests
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from ..version import __version__
 
@@ -45,7 +45,11 @@ def __init__(self, api_key, url, timeout):
         self.url = url
         self.timeout = timeout
 
-    def check(self, prompt: str, detect: Optional[bool] = False) -> Dict[str, Any]:
+    def check(
+            self,
+            prompt: Union[str, List[str]],
+            detect: Optional[bool] = False
+    ) -> Dict[str, Any]:
         """
         Creates a prompt injection check request in the Prediction Guard /injection API.
 

predictionguard/src/pii.py

@@ -1,7 +1,7 @@
 import json
 
 import requests
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from ..version import __version__
 
@@ -47,20 +47,26 @@ def __init__(self, api_key, url, timeout):
         self.timeout = timeout
 
     def check(
-        self, prompt: str, replace: bool, replace_method: Optional[str] = "random"
+        self,
+        prompt: Union[str, List[str]],
+        replace: bool,
+        replace_method: str = "random",
+        entity_list: Optional[list] = None
     ) -> Dict[str, Any]:
         """Creates a PII checking request for the Prediction Guard /PII API.
 
-        :param text: The text to check for PII.
+        :param prompt: The prompt to check for PII.
         :param replace: Whether to replace PII if it is present.
         :param replace_method: Method to replace PII if it is present.
+        :param entity_list: List of entities for the PII check to ignore.
+        :return: The PII checking response.
         """
 
         # Run _check_pii
-        choices = self._check_pii(prompt, replace, replace_method)
+        choices = self._check_pii(prompt, replace, replace_method, entity_list)
         return choices
 
-    def _check_pii(self, prompt, replace, replace_method):
+    def _check_pii(self, prompt, replace, replace_method, entity_list):
         """Function to check for PII."""
 
         headers = {
@@ -73,6 +79,7 @@ def _check_pii(self, prompt, replace, replace_method):
             "prompt": prompt,
             "replace": replace,
             "replace_method": replace_method,
+            "entity_list": entity_list
         }
 
         payload = json.dumps(payload_dict)

predictionguard/version.py

@@ -1,2 +1,2 @@
 # Setting the package version
-__version__ = "2.9.2"
+__version__ = "2.9.3"

pyproject.toml

@@ -1,25 +1,28 @@
 [build-system]
-requires = ["flit_core >=3.2,<4"]
-build-backend = "flit_core.buildapi"
+requires = ["hatchling"]
+build-backend = "hatchling.build"
 
 [project]
 name = "predictionguard"
 authors = [
-    {name = "Daniel Whitenack", email = "dan@predictionguard.com"},
-    {name = "Jacob Mansdorfer", email = "jacob@predictionguard.com"}
+    {name = "Jacob Mansdorfer", email = "jacob@predictionguard.com"},
+    {name = "Daniel Whitenack", email = "dan@predictionguard.com"}
     ]
 readme = "README.md"
+requires-python = ">=3.12.8"
 license = {file = "LICENSE"}
 classifiers = ["License :: OSI Approved :: MIT License"]
 dynamic = ["version", "description"]
 description-file = "README.md"
 dependencies = [
-    "tabulate>=0.8.10",
-    "requests>=2.27.1",
+    "tabulate>=0.9.0",
+    "requests>=2.32.5",
 ]
 
 [project.optional-dependencies]
 dev = [
+    "pytest",
+    "ruff",
     "sphinx",
     "sphinx-autodoc-typehints",
     "sphinx_rtd_theme",
@@ -28,3 +31,6 @@ dev = [
 
 [project.urls]
 Home = "https://predictionguard.com"
+
+[tool.hatch.version]
+path = "predictionguard/version.py"