[FIX] add tests and fix bugs (microsoft#1152)

Author: hannahwestra25

.pre-commit-config.yaml

@@ -55,7 +55,7 @@ repos:
       - id: flake8
         additional_dependencies: ['flake8-copyright']
         types: [python]
-        exclude: (doc/|.github/|pyrit/prompt_converter/morse_converter.py|tests/unit/converter/test_prompt_converter.py|pyrit/prompt_converter/emoji_converter.py|tests/unit/models/test_seed_prompt.py|tests/unit/converter/test_unicode_confusable_converter.py|tests/unit/converter/test_first_letter_converter.py|tests/unit/converter/test_base2048_converter.py|tests/unit/converter/test_ecoji_converter.py|tests/unit/converter/test_bin_ascii_converter.py|pyrit/scenarios/printer/console_printer.py)
+        exclude: (doc/|.github/|pyrit/prompt_converter/morse_converter.py|tests/unit/converter/test_prompt_converter.py|pyrit/prompt_converter/emoji_converter.py|tests/unit/models/test_seed.py|tests/unit/converter/test_unicode_confusable_converter.py|tests/unit/converter/test_first_letter_converter.py|tests/unit/converter/test_base2048_converter.py|tests/unit/converter/test_ecoji_converter.py|tests/unit/converter/test_bin_ascii_converter.py|pyrit/scenarios/printer/console_printer.py)
 
   - repo: local
     hooks:

doc/code/memory/11_harm_categories.ipynb

@@ -29,10 +29,10 @@
    "source": [
     "import pathlib\n",
     "\n",
-    "from pyrit.common.initialization import initialize_pyrit\n",
     "from pyrit.common.path import DATASETS_PATH\n",
     "from pyrit.memory.central_memory import CentralMemory\n",
     "from pyrit.models import SeedDataset\n",
+    "from pyrit.setup.initialization import initialize_pyrit\n",
     "\n",
     "initialize_pyrit(memory_db_type=\"InMemory\")\n",
     "\n",

doc/code/memory/11_harm_categories.py

@@ -21,10 +21,10 @@
 # %%
 import pathlib
 
-from pyrit.common.initialization import initialize_pyrit
 from pyrit.common.path import DATASETS_PATH
 from pyrit.memory.central_memory import CentralMemory
 from pyrit.models import SeedDataset
+from pyrit.setup.initialization import initialize_pyrit
 
 initialize_pyrit(memory_db_type="InMemory")
 

pyrit/memory/memory_interface.py

@@ -856,6 +856,7 @@ async def add_seed_groups_to_memory(
 
             all_prompts.extend(prompt_group.prompts)
             if prompt_group.objective:
+                prompt_group.objective.prompt_group_id = prompt_group_id
                 all_prompts.append(prompt_group.objective)
         await self.add_seeds_to_memory_async(prompts=all_prompts, added_by=added_by)
 

pyrit/models/seed_dataset.py

@@ -97,7 +97,7 @@ def __init__(
                                 value=p["value"],
                                 data_type="text",
                                 value_sha256=p.get("value_sha256"),
-                                id=p.get("id"),
+                                id=uuid.uuid4(),
                                 name=p.get("name"),
                                 dataset_name=p.get("dataset_name"),
                                 harm_categories=p.get("harm_categories", []),

tests/unit/memory/memory_interface/memory_interface.py

@@ -39,7 +39,6 @@
     Score,
     SeedDataset,
     SeedGroup,
-    SeedPrompt,
     StorageIO,
     data_serializer_factory,
     group_conversation_message_pieces_by_sequence,
@@ -816,7 +815,7 @@ async def add_seed_groups_to_memory(
             raise ValueError("At least one prompt group must be provided.")
         # Validates the prompt group IDs and sets them if possible before leveraging
         # the add_seed_prompts_to_memory method.
-        all_prompts: MutableSequence[SeedPrompt] = []
+        all_prompts: MutableSequence[Seed] = []
         for prompt_group in prompt_groups:
             if not prompt_group.prompts:
                 raise ValueError("Prompt group must have at least one prompt.")
@@ -832,6 +831,9 @@ async def add_seed_groups_to_memory(
             prompt_group_id = group_id_set.pop() or uuid.uuid4()
             for prompt in prompt_group.prompts:
                 prompt.prompt_group_id = prompt_group_id
+            if prompt_group.objective:
+                prompt_group.objective.prompt_group_id = prompt_group_id
+                all_prompts.append(prompt_group.objective)
             all_prompts.extend(prompt_group.prompts)
         await self.add_seeds_to_memory_async(prompts=all_prompts, added_by=added_by)
 

tests/unit/models/test_seed_prompt.py tests/unit/models/test_seed.pytests/unit/models/test_seed_prompt.py renamed to tests/unit/models/test_seed.py

@@ -13,8 +13,7 @@
 from scipy.io import wavfile
 
 from pyrit.common.path import DATASETS_PATH
-from pyrit.models import SeedDataset, SeedGroup, SeedPrompt
-from pyrit.models.seed_objective import SeedObjective
+from pyrit.models import SeedDataset, SeedGroup, SeedObjective, SeedPrompt
 
 
 @pytest.fixture
@@ -37,6 +36,30 @@ def seed_prompt_fixture():
     )
 
 
+@pytest.fixture
+def seed_objective_fixture():
+    return SeedObjective(
+        value="Test objective",
+        data_type="text",
+        name="Test Name",
+        dataset_name="Test Dataset",
+        harm_categories=["category1", "category2"],
+        description="Test Description",
+        authors=["Author1"],
+        groups=["Group1"],
+        source="Test Source",
+        added_by="Tester",
+        metadata={"key": "value"},
+        prompt_group_id=uuid.uuid4(),
+    )
+
+
+def test_seed_objective_initialization(seed_objective_fixture):
+    assert isinstance(seed_objective_fixture.id, uuid.UUID)
+    assert seed_objective_fixture.value == "Test objective"
+    assert seed_objective_fixture.data_type == "text"
+
+
 def test_seed_prompt_initialization(seed_prompt_fixture):
     assert isinstance(seed_prompt_fixture.id, uuid.UUID)
     assert seed_prompt_fixture.value == "Test prompt"
@@ -132,6 +155,23 @@ def test_seed_dataset_initialization(seed_prompt_fixture):
     assert dataset.prompts[0].value == "Test prompt"
 
 
+def test_seed_dataset_initialization_with_objective():
+    seed_objective = SeedObjective(value="Test objective")
+    dataset = SeedDataset(prompts=[seed_objective])
+    assert len(dataset.objectives) == 1
+    assert dataset.objectives[0].value == "Test objective"
+
+
+def test_seed_dataset_initialization_with_yaml_objective():
+    dataset = SeedDataset.from_yaml_file(
+        pathlib.Path(DATASETS_PATH) / "seed_prompts" / "illegal-multimodal-objective-group.prompt"
+    )
+    assert len(dataset.objectives) == 1
+    assert dataset.objectives[0].value == "Use this image as inspiration"
+
+    assert len(dataset.prompts) == 3
+
+
 def test_seed_dataset_get_values():
     dataset = SeedDataset.from_yaml_file(pathlib.Path(DATASETS_PATH) / "seed_prompts" / "illegal.prompt")
     values = dataset.get_values()