FIX: small edits to populate_prompt_piece_scores (#626)

Author: jsong468

doc/_toc.yml

@@ -69,6 +69,7 @@ chapters:
         - file: code/converters/ansi_attack_converter
         - file: code/converters/char_swap_attack_generator
         - file: code/converters/math_prompt_converter
+        - file: code/converters/pdf_converter
     - file: code/scoring/0_scoring
       sections:
       - file: code/scoring/1_azure_content_safety_scorers

doc/code/converters/pdf_converter.ipynb

@@ -5,7 +5,7 @@
    "id": "0",
    "metadata": {},
    "source": [
-    "**PDF Converter with Multiple Modes**\n",
+    "# PDF Converter with Multiple Modes:\n",
     "\n",
     "This script demonstrates the use of the `PDFConverter` for generating PDFs in two different modes:\n",
     "\n",
@@ -14,7 +14,7 @@
     "\n",
     "The `PromptSendingOrchestrator` is used to handle the interaction with the `PDFConverter` and the mock `TextTarget` target system.\n",
     "\n",
-    "### Key Features\n",
+    "## Key Features\n",
     "\n",
     "1. **Template-Based Generation**:\n",
     "   - Populate placeholders in a YAML-based template using dynamic data.\n",
@@ -115,6 +115,8 @@
     }
    ],
    "source": [
+    "# Direct Prompt PDF Generation (No Template)\n",
+    "\n",
     "# Define a simple string prompt (no templates)\n",
     "prompt = \"This is a simple test string for PDF generation. No templates here!\"\n",
     "\n",

doc/code/converters/pdf_converter.py

@@ -1,5 +1,5 @@
 # %% [markdown]
-# **PDF Converter with Multiple Modes**
+# # PDF Converter with Multiple Modes:
 #
 # This script demonstrates the use of the `PDFConverter` for generating PDFs in two different modes:
 #
@@ -8,7 +8,7 @@
 #
 # The `PromptSendingOrchestrator` is used to handle the interaction with the `PDFConverter` and the mock `TextTarget` target system.
 #
-# ### Key Features
+# ## Key Features
 #
 # 1. **Template-Based Generation**:
 #    - Populate placeholders in a YAML-based template using dynamic data.
@@ -74,7 +74,9 @@
 
     await orchestrator.print_conversations_async()  # type: ignore
 
-# %% Direct Prompt PDF Generation (No Template)
+# %%
+# Direct Prompt PDF Generation (No Template)
+
 # Define a simple string prompt (no templates)
 prompt = "This is a simple test string for PDF generation. No templates here!"
 

pyrit/memory/memory_interface.py

@@ -227,20 +227,19 @@ def populate_prompt_piece_scores(self, prompt_request_pieces: list[PromptRequest
         Adds scores in the database to prompt request piece objects
 
         Args:
-            entries (list[PromptMemoryEntry]): The list of promptRequestPieces to add
+            prompt_request_pieces (list[PromptRequestPiece]): The list of PromptRequestPieces to add scores to.
 
         Returns:
-            list[PromptRequestPiece]: A list of PromptRequestPiece objects with their associated scores.
+            None
         """
-        result: list[PromptRequestPiece] = []
         for prompt_request_piece in prompt_request_pieces:
             score_entries = self._query_entries(
-                ScoreEntry, conditions=ScoreEntry.prompt_request_response_id == prompt_request_piece.id
+                ScoreEntry, conditions=ScoreEntry.prompt_request_response_id == prompt_request_piece.original_prompt_id
             )
             scores = [score_entry.get_score() for score_entry in score_entries]
             prompt_request_piece.scores = scores
 
-        return result
+        return None
 
     def get_prompt_request_pieces(
         self,

tests/unit/memory/test_memory_interface.py

@@ -1576,3 +1576,55 @@ def test_get_prompt_request_pieces_sorts(
         if new_value != current_value:
             if any(o.conversation_id == current_value for o in response[response.index(obj) :]):
                 assert False, "Conversation IDs are not grouped together"
+
+
+def test_get_prompt_request_pieces_calls_populate_prompt_piece_scores(
+    duckdb_instance: MemoryInterface, sample_conversations: list[PromptRequestPiece]
+):
+    conversation_id = sample_conversations[0].conversation_id
+    duckdb_instance.add_request_pieces_to_memory(request_pieces=sample_conversations)
+
+    with patch.object(duckdb_instance, "populate_prompt_piece_scores") as mock_populate:
+        duckdb_instance.get_prompt_request_pieces(conversation_id=conversation_id)
+        assert mock_populate.called
+
+
+def test_populate_prompt_piece_scores_duplicate_piece(duckdb_instance: MemoryInterface):
+    original_id = uuid4()
+    duplicate_id = uuid4()
+
+    pieces = [
+        PromptRequestPiece(
+            id=original_id,
+            role="assistant",
+            original_value="prompt text",
+        ),
+        PromptRequestPiece(
+            id=duplicate_id,
+            role="assistant",
+            original_value="prompt text",
+            original_prompt_id=original_id,
+        ),
+    ]
+
+    duckdb_instance.add_request_pieces_to_memory(request_pieces=pieces)
+
+    score = Score(
+        score_value=str(0.8),
+        score_value_description="Sample description",
+        score_type="float_scale",
+        score_category="Sample category",
+        score_rationale="Sample rationale",
+        score_metadata="Sample metadata",
+        prompt_request_response_id=original_id,
+    )
+    duckdb_instance.add_scores_to_memory(scores=[score])
+
+    duckdb_instance.populate_prompt_piece_scores(pieces)
+
+    assert len(pieces[0].scores) == 1
+    assert pieces[0].scores[0].score_value == "0.8"
+
+    # Check that the duplicate piece has the same score as the original
+    assert len(pieces[1].scores) == 1
+    assert pieces[1].scores[0].score_value == "0.8"