Include datasets in package (Azure#63)

MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include pyrit *.yaml
+recursive-include pyrit *.prompt

doc/code/aml_endpoints.ipynb

@@ -122,6 +122,7 @@
     "        pathlib.Path(os.getcwd())\n",
     "        / \"..\"\n",
     "        / \"..\"\n",
+    "        / \"pyrit\"\n",
     "        / \"datasets\"\n",
     "        / \"attack_strategies\"\n",
     "        / \"multi_turn_chat\"\n",

doc/code/aml_endpoints.py

@@ -69,6 +69,7 @@
         pathlib.Path(os.getcwd())
         / ".."
         / ".."
+        / "pyrit"
         / "datasets"
         / "attack_strategies"
         / "multi_turn_chat"

doc/demo/1_gandalf.ipynb

@@ -108,6 +108,7 @@
     "        pathlib.Path(os.getcwd())\n",
     "        / \"..\"\n",
     "        / \"..\"\n",
+    "        / \"pyrit\"\n",
     "        / \"datasets\"\n",
     "        / \"attack_strategies\"\n",
     "        / \"multi_turn_chat\"\n",

doc/demo/1_gandalf.py

@@ -76,6 +76,7 @@
         pathlib.Path(os.getcwd())
         / ".."
         / ".."
+        / "pyrit"
         / "datasets"
         / "attack_strategies"
         / "multi_turn_chat"

doc/how_to_guide.ipynb

@@ -238,7 +238,7 @@
    "source": [
     "This way, the red teamer can create lots of different conversations based on a relatively small number\n",
     "of templates. For better maintainability, we suggest storing the prompt templates in YAML files (see,\n",
-    "for example, `datasets/attack_strategies/multi-turn-chat/red_team_chatbot_with_objective.yaml`).\n",
+    "for example, `pyrit/datasets/attack_strategies/multi-turn-chat/red_team_chatbot_with_objective.yaml`).\n",
     "\n",
     "PyRIT offers various integration choices for the `RedTeamingBot`, including\n",
     "[Azure ML managed online endpoints](../doc/code/aml_endpoints.ipynb),\n",

doc/how_to_guide.py

@@ -124,7 +124,7 @@
 # %% [markdown]
 # This way, the red teamer can create lots of different conversations based on a relatively small number
 # of templates. For better maintainability, we suggest storing the prompt templates in YAML files (see,
-# for example, `datasets/attack_strategies/multi-turn-chat/red_team_chatbot_with_objective.yaml`).
+# for example, `pyrit/datasets/attack_strategies/multi-turn-chat/red_team_chatbot_with_objective.yaml`).
 #
 # PyRIT offers various integration choices for the `RedTeamingBot`, including
 # [Azure ML managed online endpoints](../doc/code/aml_endpoints.ipynb),

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pyrit"
-version = "0.1.0"
+version = "0.1.2.dev0"
 description = "The Python Risk Identification Tool for LLMs (PyRIT) is a library used to assess the robustness of LLMs"
 authors = [
     { name = "Microsoft AI Red Team", email = "[email protected]" },
@@ -83,11 +83,11 @@ target-version = ["py310"]
 line-length = 120
 
 [build-system]
-requires = ["setuptools", "setuptools-scm"]
+requires = ["setuptools", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
-include = ["pyrit", "pyrit.*", "datasets", "datasets.*"]
+include = ["pyrit", "pyrit.*"]
 
 [tool.jupytext]
 formats = "ipynb,py:percent"

pyrit/__init__.py

@@ -6,4 +6,4 @@
 
 __name__ = "pyrit"
 # Remove dev suffix when releasing and keep in sync with pyproject.toml
-__version__ = "0.1.1.dev0"
+__version__ = "0.1.2.dev0"

pyrit/agent/red_teaming_bot.py

@@ -7,7 +7,7 @@
 from pyrit.interfaces import ChatSupport
 from pyrit.memory import FileMemory, MemoryInterface
 from pyrit.models import ChatMessage, PromptTemplate
-from pyrit.common.path import HOME_PATH
+from pyrit.common.path import DATASETS_PATH
 
 RED_TEAM_CHATBOT_ROLE_PREFIX = "airt-bot-uuid"
 RED_TEAM_CONVERSATION_END_TOKEN = "<|done|>"
@@ -45,8 +45,7 @@ def __init__(
             attack_strategy
             if attack_strategy
             else PromptTemplate.from_yaml_file(
-                pathlib.Path(HOME_PATH)
-                / "datasets"
+                pathlib.Path(DATASETS_PATH)
                 / "attack_strategies"
                 / "multi_turn_chat"
                 / "red_team_chatbot_with_objective.yaml"

pyrit/common/path.py

@@ -11,7 +11,7 @@
 LOG_PATH = pathlib.Path(PYRIT_PATH, "..", "results", "logs.txt").resolve()
 LOG_PATH.touch(exist_ok=True)
 
-DATASETS_PATH = pathlib.Path(PYRIT_PATH, "..", "datasets").resolve()
+DATASETS_PATH = pathlib.Path(PYRIT_PATH, "datasets").resolve()
 
 # Points to the root of the project
 HOME_PATH = pathlib.Path(PYRIT_PATH, "..").resolve()

tests/test_red_teaming_bot.py

@@ -13,7 +13,7 @@
 from pyrit.chat import AzureOpenAIChat
 from pyrit.models import PromptTemplate
 from pyrit.memory import FileMemory
-from pyrit.common.path import HOME_PATH
+from pyrit.common.path import DATASETS_PATH
 
 
 @pytest.fixture
@@ -42,11 +42,7 @@ def chat_completion_engine() -> AzureOpenAIChat:
 @pytest.fixture
 def red_teaming_bot(chat_completion_engine: AzureOpenAIChat, tmp_path: pathlib.Path):
     attack_strategy = PromptTemplate.from_yaml_file(
-        pathlib.Path(HOME_PATH)
-        / "datasets"
-        / "attack_strategies"
-        / "multi_turn_chat"
-        / "red_team_chatbot_with_objective.yaml"
+        pathlib.Path(DATASETS_PATH) / "attack_strategies" / "multi_turn_chat" / "red_team_chatbot_with_objective.yaml"
     )
 
     file_memory = FileMemory(filepath=tmp_path / "test.json.memory")