MAINT: Adding Identifiers to Memory (Azure#145)

* Adding Identifiers to Memory * Adding squash commit shortcut * reverting notebook * PR feedback * PR feedback
uskr · Apr 18, 2024 · 2775b76 · 2775b76
1 parent 4c649a1
commit 2775b76
Show file tree

Hide file tree

Showing 54 changed files with 901 additions and 1,368 deletions.
diff --git a/.github/squash_commits.ps1 b/.github/squash_commits.ps1
@@ -0,0 +1,29 @@
+param(
+    [string]$CommitMessage,
+    [string]$fork
+)
+
+$mainBranch = 'main'
+$featureBranch = git rev-parse --abbrev-ref HEAD
+
+git fetch origin $mainBranch
+
+# Check out the feature branch (though you should already be on it)
+git checkout $featureBranch
+
+git rebase -i origin/$mainBranch
+
+# Squash commits
+# Note: The interactive rebase '-i' will open an editor to squash commits manually
+# Replace 'pick' with 'squash' beside all but the first commit to combine them
+
+# If you're not comfortable with the interactive mode or want to automate:
+# Assuming you want to squash all commits made on the feature branch since it diverged from main:
+$commitCount = (git rev-list --count HEAD ^origin/$mainBranch)
+if ($commitCount -gt 1) {
+    git reset --soft "HEAD~$commitCount"
+    git commit -m $CommitMessage
+}
+
+# Push changes to the remote repository
+git push $fork $featureBranch --force
diff --git a/doc/code/memory/memory.ipynb b/doc/code/memory/memory.ipynb
@@ -2,33 +2,93 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "8015157a",
+   "id": "da395b00",
    "metadata": {},
    "source": [
-    "The `pyrit.memory` module provides functionality to keep track of the conversation history. In a nutshell, this can be used as follows"
+    "The memory module is the primary way pyrit keeps track of requests and responses to targets. The schema is found in `memory_models.py` and can be programatically viewed as follows\n"
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "3362c88f",
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a2a45a7c",
    "metadata": {
-    "lines_to_next_cell": 0
+    "execution": {
+     "iopub.execute_input": "2024-04-17T21:46:23.697329Z",
+     "iopub.status.busy": "2024-04-17T21:46:23.696330Z",
+     "iopub.status.idle": "2024-04-17T21:46:25.974202Z",
+     "shell.execute_reply": "2024-04-17T21:46:25.974202Z"
+    }
    },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\rlundeen\\AppData\\Local\\anaconda3\\envs\\pyrit-dev\\lib\\site-packages\\duckdb_engine\\__init__.py:565: SAWarning: Did not recognize type 'list' of column 'embedding'\n",
+      "  columns = self._get_columns_info(rows, domains, enums, schema)  # type: ignore[attr-defined]\n",
+      "c:\\Users\\rlundeen\\AppData\\Local\\anaconda3\\envs\\pyrit-dev\\lib\\site-packages\\duckdb_engine\\__init__.py:180: DuckDBEngineWarning: duckdb-engine doesn't yet support reflection on indices\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Schema for EmbeddingData:\n",
+      "  Column id (UUID)\n",
+      "  Column embedding (NULL)\n",
+      "  Column embedding_type_name (VARCHAR)\n",
+      "Schema for PromptMemoryEntries:\n",
+      "  Column id (UUID)\n",
+      "  Column role (VARCHAR)\n",
+      "  Column conversation_id (VARCHAR)\n",
+      "  Column sequence (INTEGER)\n",
+      "  Column timestamp (TIMESTAMP)\n",
+      "  Column labels (VARCHAR)\n",
+      "  Column prompt_metadata (VARCHAR)\n",
+      "  Column converter_identifiers (VARCHAR)\n",
+      "  Column prompt_target_identifier (VARCHAR)\n",
+      "  Column orchestrator_identifier (VARCHAR)\n",
+      "  Column response_error (VARCHAR)\n",
+      "  Column original_prompt_data_type (VARCHAR)\n",
+      "  Column original_prompt_text (VARCHAR)\n",
+      "  Column original_prompt_data_sha256 (VARCHAR)\n",
+      "  Column converted_prompt_data_type (VARCHAR)\n",
+      "  Column converted_prompt_text (VARCHAR)\n",
+      "  Column converted_prompt_data_sha256 (VARCHAR)\n"
+     ]
+    }
+   ],
    "source": [
-    "The PyRIT DuckDB database comprises of tables defined in the `pyrit.memory.memory_models` module.\n",
-    "\n"
+    "# Copyright (c) Microsoft Corporation.\n",
+    "# Licensed under the MIT license.\n",
+    "\n",
+    "from pyrit.memory import DuckDBMemory\n",
+    "memory = DuckDBMemory()\n",
+    "\n",
+    "memory.print_schema()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "30322f05",
+   "metadata": {},
+   "source": [
+    "\n",
+    "The `pyrit.memory` module provides functionality to keep track of the conversation history. In a nutshell, this can be used as follows"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
-   "id": "7fe509c1",
+   "execution_count": 2,
+   "id": "b58266be",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-04-15T21:31:23.317029Z",
-     "iopub.status.busy": "2024-04-15T21:31:23.317029Z",
-     "iopub.status.idle": "2024-04-15T21:31:24.597203Z",
-     "shell.execute_reply": "2024-04-15T21:31:24.596522Z"
+     "iopub.execute_input": "2024-04-17T21:46:25.974202Z",
+     "iopub.status.busy": "2024-04-17T21:46:25.974202Z",
+     "iopub.status.idle": "2024-04-17T21:46:26.044580Z",
+     "shell.execute_reply": "2024-04-17T21:46:26.043420Z"
     },
     "lines_to_next_cell": 2
    },
@@ -44,16 +104,13 @@
     }
    ],
    "source": [
-    "# Copyright (c) Microsoft Corporation.\n",
-    "# Licensed under the MIT license.\n",
+    "\n",
     "\n",
     "from uuid import uuid4\n",
-    "from pyrit.memory import DuckDBMemory\n",
     "from pyrit.models import PromptRequestPiece\n",
     "\n",
     "conversation_id = str(uuid4())\n",
     "\n",
-    "memory = DuckDBMemory()\n",
     "message_list = [\n",
     "    PromptRequestPiece(\n",
     "        role=\"user\", original_prompt_text=\"Hi, chat bot! This is my initial prompt.\", conversation_id=conversation_id\n",
@@ -79,24 +136,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "id": "71d7bc29",
+   "execution_count": 3,
+   "id": "31c2b09a",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-04-15T21:31:24.598033Z",
-     "iopub.status.busy": "2024-04-15T21:31:24.598033Z",
-     "iopub.status.idle": "2024-04-15T21:31:24.647986Z",
-     "shell.execute_reply": "2024-04-15T21:31:24.646493Z"
+     "iopub.execute_input": "2024-04-17T21:46:26.048616Z",
+     "iopub.status.busy": "2024-04-17T21:46:26.047620Z",
+     "iopub.status.idle": "2024-04-17T21:46:26.097887Z",
+     "shell.execute_reply": "2024-04-17T21:46:26.096888Z"
     }
    },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "user: this is updated field\n",
-      "assistant: this is updated field\n",
-      "user: this is updated field\n"
+      "None: user: this is updated field\n",
+      "None: assistant: this is updated field\n",
+      "None: user: this is updated field\n"
      ]
     }
    ],
@@ -115,14 +172,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "id": "8df85019",
+   "execution_count": 4,
+   "id": "78c2f09e",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-04-15T21:31:24.647986Z",
-     "iopub.status.busy": "2024-04-15T21:31:24.647986Z",
-     "iopub.status.idle": "2024-04-15T21:31:24.680328Z",
-     "shell.execute_reply": "2024-04-15T21:31:24.679382Z"
+     "iopub.execute_input": "2024-04-17T21:46:26.102395Z",
+     "iopub.status.busy": "2024-04-17T21:46:26.100889Z",
+     "iopub.status.idle": "2024-04-17T21:46:26.128226Z",
+     "shell.execute_reply": "2024-04-17T21:46:26.127219Z"
     }
    },
    "outputs": [],
@@ -134,7 +191,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7093070e",
+   "id": "e5110a2b",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -145,9 +202,9 @@
    "cell_metadata_filter": "-all"
   },
   "kernelspec": {
-   "display_name": "pyrit_kernel",
+   "display_name": "pyrit-dev",
    "language": "python",
-   "name": "pyrit_kernel"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

diff --git a/doc/code/memory/memory.py b/doc/code/memory/memory.py
@@ -1,21 +1,29 @@
 # %% [markdown]
-# The `pyrit.memory` module provides functionality to keep track of the conversation history. In a nutshell, this can be used as follows
-
-# %% [markdown]
-# The PyRIT DuckDB database comprises of tables defined in the `pyrit.memory.memory_models` module.
-#
+# The memory module is the primary way pyrit keeps track of requests and responses to targets. The schema is found in `memory_models.py` and can be programatically viewed as follows
 #
+
 # %%
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from uuid import uuid4
 from pyrit.memory import DuckDBMemory
+
+memory = DuckDBMemory()
+
+memory.print_schema()
+
+# %% [markdown]
+#
+# The `pyrit.memory` module provides functionality to keep track of the conversation history. In a nutshell, this can be used as follows:
+
+# %%
+
+
+from uuid import uuid4
 from pyrit.models import PromptRequestPiece
 
 conversation_id = str(uuid4())
 
-memory = DuckDBMemory()
 message_list = [
     PromptRequestPiece(
         role="user", original_prompt_text="Hi, chat bot! This is my initial prompt.", conversation_id=conversation_id

diff --git a/doc/code/memory/memory_export_to_json.ipynb b/doc/code/memory/memory_export_to_json.ipynb
@@ -37,7 +37,7 @@
     "\n",
     "from pyrit.memory import DuckDBMemory\n",
     "from uuid import uuid4\n",
-    "from pyrit.models.prompt_request_piece import PromptRequestPiece\n",
+    "from pyrit.models import PromptRequestPiece\n",
     "\n",
     "duckdb_memory = DuckDBMemory()\n",
     "duckdb_memory.export_all_tables()"
@@ -136,9 +136,9 @@
    "cell_metadata_filter": "-all"
   },
   "kernelspec": {
-   "display_name": "pyrit_kernel",
+   "display_name": "pyrit-dev",
    "language": "python",
-   "name": "pyrit_kernel"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

diff --git a/doc/code/memory/memory_export_to_json.py b/doc/code/memory/memory_export_to_json.py
@@ -12,7 +12,7 @@
 
 from pyrit.memory import DuckDBMemory
 from uuid import uuid4
-from pyrit.models.prompt_request_piece import PromptRequestPiece
+from pyrit.models import PromptRequestPiece
 
 duckdb_memory = DuckDBMemory()
 duckdb_memory.export_all_tables()

diff --git a/doc/code/targets/azure_ml_chat.ipynb b/doc/code/targets/azure_ml_chat.ipynb
@@ -57,7 +57,7 @@
     "# Licensed under the MIT license.\n",
     "\n",
     "from pyrit.common import default_values\n",
-    "from pyrit.models.prompt_request_piece import PromptRequestPiece\n",
+    "from pyrit.models import PromptRequestPiece\n",
     "from pyrit.prompt_target import AzureMLChatTarget\n",
     "\n",
     "\n",

diff --git a/doc/code/targets/azure_ml_chat.py b/doc/code/targets/azure_ml_chat.py
@@ -27,7 +27,7 @@
 # Licensed under the MIT license.
 
 from pyrit.common import default_values
-from pyrit.models.prompt_request_piece import PromptRequestPiece
+from pyrit.models import PromptRequestPiece
 from pyrit.prompt_target import AzureMLChatTarget
 
 

diff --git a/doc/code/targets/dall_e_target.ipynb b/doc/code/targets/dall_e_target.ipynb
@@ -30,8 +30,8 @@
     "import uuid\n",
     "\n",
     "from pyrit.common import default_values\n",
-    "from pyrit.models.prompt_request_piece import PromptRequestPiece\n",
-    "from pyrit.prompt_target import DallETarget\n",
+    "from pyrit.models import PromptRequestPiece\n",
+    "from pyrit.prompt_target import DALLETarget\n",
     "from pyrit.prompt_target.dall_e_target import ResponseFormat\n",
     "\n",
     "test_conversation_id = str(uuid.uuid4())\n",
@@ -40,7 +40,7 @@
     "\n",
     "default_values.load_default_env()\n",
     "\n",
-    "img_prompt_target = DallETarget(\n",
+    "img_prompt_target = DALLETarget(\n",
     "    deployment_name=os.environ.get(\"AZURE_DALLE_DEPLOYMENT\"),\n",
     "    endpoint=os.environ.get(\"AZURE_DALLE_ENDPOINT\"),\n",
     "    api_key=os.environ.get(\"AZURE_DALLE_API_KEY\"),\n",

diff --git a/doc/code/targets/dall_e_target.py b/doc/code/targets/dall_e_target.py
@@ -10,7 +10,7 @@
 import uuid
 
 from pyrit.common import default_values
-from pyrit.models.prompt_request_piece import PromptRequestPiece
+from pyrit.models import PromptRequestPiece
 from pyrit.prompt_target import DALLETarget
 from pyrit.prompt_target.dall_e_target import ResponseFormat
 

diff --git a/doc/code/targets/prompt_targets.ipynb b/doc/code/targets/prompt_targets.ipynb
@@ -45,7 +45,6 @@
     "\n",
     "from pyrit.models import PromptTemplate\n",
     "from pyrit.common.path import DATASETS_PATH\n",
-    "from pyrit.models.prompt_request_piece import PromptRequestPiece\n",
     "\n",
     "jailbreak_template = PromptTemplate.from_yaml_file(\n",
     "    pathlib.Path(DATASETS_PATH) / \"prompt_templates\" / \"jailbreak\" / \"jailbreak_1.yaml\"\n",
@@ -88,7 +87,7 @@
     }
    ],
    "source": [
-    "from pyrit.models.prompt_request_piece import PromptRequestPiece\n",
+    "from pyrit.models import PromptRequestPiece\n",
     "from pyrit.prompt_target import AzureOpenAIChatTarget\n",
     "from pyrit.common import default_values\n",
     "\n",
@@ -171,9 +170,9 @@
    "cell_metadata_filter": "-all"
   },
   "kernelspec": {
-   "display_name": "pyrit_kernel",
+   "display_name": "pyrit-dev",
    "language": "python",
-   "name": "pyrit_kernel"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

diff --git a/doc/code/targets/prompt_targets.py b/doc/code/targets/prompt_targets.py
@@ -16,7 +16,7 @@
 
 from pyrit.models import PromptTemplate
 from pyrit.common.path import DATASETS_PATH
-from pyrit.models.prompt_request_piece import PromptRequestPiece
+from pyrit.models import PromptRequestPiece
 
 jailbreak_template = PromptTemplate.from_yaml_file(
     pathlib.Path(DATASETS_PATH) / "prompt_templates" / "jailbreak" / "jailbreak_1.yaml"
@@ -31,7 +31,7 @@
 # _Note:_ to run this section of the demo you need to setup your `.env` file to properly authenticate to an Azure OpenAI endpoint as described [here](../../setup/setup_azure.md).
 
 # %%
-from pyrit.models.prompt_request_piece import PromptRequestPiece
+from pyrit.models import PromptRequestPiece
 from pyrit.prompt_target import AzureOpenAIChatTarget
 from pyrit.common import default_values