FEAT Add Scores to Data Export with PromptRequestPiece data (#617)

Author: eugeniavkim

doc/code/memory/9_exporting_data.ipynb

@@ -6,41 +6,126 @@
 #       format_name: percent
 #       format_version: '1.3'
 #       jupytext_version: 1.16.4
+#   kernelspec:
+#     display_name: pyrit-kernel
+#     language: python
+#     name: pyrit-kernel
 # ---
 
 # %% [markdown]
 # # 9. Exporting Data Example
 #
-# This notebook shows all the different ways to export data from memory. This first example exports all conversations from Azure SQL memory with their respective score values in a JSON format. Without using the database query editor, these export functions allow for a quick approach to exporting data from memory.
+# This notebook shows different ways to export data from memory. This first example exports all conversations from local DuckDB memory with their respective score values in a JSON format. The data can currently be exported both as JSON file or a CSV file that will be saved in your results folder within PyRIT. The CSV export is commented out below. In this example, all conversations are exported, but by using other export functions from `memory_interface`, we can export by specific labels and other methods.
 
 # %%
-from pyrit.memory.azure_sql_memory import AzureSQLMemory
+from uuid import uuid4
+
 from pyrit.common import default_values
-from pathlib import Path
+from pyrit.common.path import RESULTS_PATH
+from pyrit.memory import DuckDBMemory, CentralMemory
+from pyrit.models import PromptRequestPiece, PromptRequestResponse
 
 default_values.load_environment_files()
 
-memory = AzureSQLMemory()
+memory = DuckDBMemory()
+CentralMemory.set_memory_instance(memory)
+
+conversation_id = str(uuid4())
+
+print(conversation_id)
+
+message_list = [
+    PromptRequestPiece(
+        role="user", original_value="Hi, chat bot! This is my initial prompt.", conversation_id=conversation_id
+    ),
+    PromptRequestPiece(
+        role="assistant", original_value="Nice to meet you! This is my response.", conversation_id=conversation_id
+    ),
+    PromptRequestPiece(
+        role="user",
+        original_value="Wonderful! This is my second prompt to the chat bot!",
+        conversation_id=conversation_id,
+    ),
+]
+
+memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[0]]))
+memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[1]]))
+memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[2]]))
+
+entries = memory.get_conversation(conversation_id=conversation_id)
+
+for entry in entries:
+    print(entry)
 
 # Define file path for export
-json_file_path = Path("conversation_and_scores_json_example")
-# csv_file_path = Path(\"conversation_and_scores_csv_example\")
+json_file_path = RESULTS_PATH / "conversation_and_scores_json_example.json"
+# csv_file_path = RESULTS_PATH / "conversation_and_scores_csv_example.csv"
 
-# Export the data to a JSON file
-conversation_with_scores = memory.export_all_conversations_with_scores(file_path=json_file_path, export_type="json")
+# # Export the data to a JSON file
+conversation_with_scores = memory.export_all_conversations(file_path=json_file_path, export_type="json")
 print(f"Exported conversation with scores to JSON: {json_file_path}")
 
 # Export the data to a CSV file
-# conversation_with_scores = memory.export_all_conversations_with_scores(file_path=json_file_path, export_type="csv")
+# conversation_with_scores = memory.export_all_conversations(file_path=csv_file_path, export_type="csv")
 # print(f"Exported conversation with scores to CSV: {csv_file_path}")
 
 # %% [markdown]
-# ## Importing Data as NumPy DataFrame
-#
-# You can use the exported JSON or CSV files to import the data as a NumPy DataFrame. This can be useful for various data manipulation and analysis tasks.
+# You can also use the exported JSON or CSV files to import the data as a NumPy DataFrame. This can be useful for various data manipulation and analysis tasks.
 
 # %%
 import pandas as pd  # type: ignore
 
 df = pd.read_json(json_file_path)
-df.head()
+df.head(1)
+
+# %% [markdown]
+# Next, we can export data from our Azure SQL database. In this example, we export the data by `conversation_id` and to a CSV file.
+
+# %%
+from pyrit.memory import AzureSQLMemory
+
+conversation_id = str(uuid4())
+
+message_list = [
+    PromptRequestPiece(
+        role="user", original_value="Hi, chat bot! This is my initial prompt.", conversation_id=conversation_id
+    ),
+    PromptRequestPiece(
+        role="assistant", original_value="Nice to meet you! This is my response.", conversation_id=conversation_id
+    ),
+    PromptRequestPiece(
+        role="user",
+        original_value="Wonderful! This is my second prompt to the chat bot!",
+        conversation_id=conversation_id,
+    ),
+]
+
+azure_memory = AzureSQLMemory()
+CentralMemory.set_memory_instance(azure_memory)
+
+azure_memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[0]]))
+azure_memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[1]]))
+azure_memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[2]]))
+
+
+entries = azure_memory.get_conversation(conversation_id=conversation_id)
+
+for entry in entries:
+    print(entry)
+
+# Define file path for export
+# json_file_path = RESULTS_PATH / "conversation_and_scores_json_example.json"
+csv_file_path = RESULTS_PATH / "conversation_and_scores_csv_example.csv"
+
+# Export the data to a JSON file
+# conversation_with_scores = azure_memory.export_conversation_by_id(conversation_id=conversation_id, file_path=json_file_path, export_type="json")
+# print(f"Exported conversation with scores to JSON: {json_file_path}")
+
+# Export the data to a CSV file
+conversation_with_scores = azure_memory.export_conversation_by_id(
+    conversation_id=conversation_id, file_path=json_file_path, export_type="csv"
+)
+print(f"Exported conversation with scores to CSV: {csv_file_path}")
+
+# Cleanup memory resources
+azure_memory.dispose_engine()

doc/code/memory/9_exporting_data.py

@@ -234,8 +234,8 @@ def _get_prompt_pieces_with_conversation_id(self, *, conversation_id: str) -> li
                 conditions=PromptMemoryEntry.conversation_id == str(conversation_id),
             )  # type: ignore
 
-            prompt_pieces: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
-            return prompt_pieces
+            result = self.get_prompt_request_pieces_with_scores(entries)
+            return result
 
         except Exception as e:
             logger.exception(f"Failed to retrieve conversation_id {conversation_id} with error {e}")
@@ -267,8 +267,9 @@ def get_all_prompt_pieces(self) -> list[PromptRequestPiece]:
         """
         Fetches all entries from the specified table and returns them as model instances.
         """
-        result: list[PromptMemoryEntry] = self.query_entries(PromptMemoryEntry)
-        return [entry.get_prompt_request_piece() for entry in result]
+        entries = self.query_entries(PromptMemoryEntry)
+        result = self.get_prompt_request_pieces_with_scores(entries)
+        return result
 
     def get_prompt_request_pieces_by_id(self, *, prompt_ids: list[str]) -> list[PromptRequestPiece]:
         """
@@ -285,7 +286,7 @@ def get_prompt_request_pieces_by_id(self, *, prompt_ids: list[str]) -> list[Prom
                 PromptMemoryEntry,
                 conditions=PromptMemoryEntry.id.in_(prompt_ids),
             )
-            result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+            result = self.get_prompt_request_pieces_with_scores(entries)
             return result
         except Exception as e:
             logger.exception(
@@ -319,7 +320,7 @@ def get_prompt_request_piece_by_memory_labels(
             sql_condition = text(conditions).bindparams(**{key: str(value) for key, value in memory_labels.items()})
 
             entries = self.query_entries(PromptMemoryEntry, conditions=sql_condition)
-            result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+            result = self.get_prompt_request_pieces_with_scores(entries)
             return result
         except Exception as e:
             logger.exception(

pyrit/memory/azure_sql_memory.py

@@ -89,7 +89,7 @@ def get_all_prompt_pieces(self) -> list[PromptRequestPiece]:
         Fetches all entries from the specified table and returns them as model instances.
         """
         entries = self.query_entries(PromptMemoryEntry)
-        result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+        result = self.get_prompt_request_pieces_with_scores(entries)
         return result
 
     def get_all_embeddings(self) -> list[EmbeddingDataEntry]:
@@ -113,8 +113,8 @@ def _get_prompt_pieces_with_conversation_id(self, *, conversation_id: str) -> li
             entries = self.query_entries(
                 PromptMemoryEntry, conditions=PromptMemoryEntry.conversation_id == str(conversation_id)
             )
-            prompt_pieces: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
-            return prompt_pieces
+            result = self.get_prompt_request_pieces_with_scores(entries)
+            return result
         except Exception as e:
             logger.exception(f"Failed to retrieve conversation_id {conversation_id} with error {e}")
             return []
@@ -134,7 +134,7 @@ def get_prompt_request_pieces_by_id(self, *, prompt_ids: list[str]) -> list[Prom
                 PromptMemoryEntry,
                 conditions=PromptMemoryEntry.id.in_(prompt_ids),
             )
-            result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+            result = self.get_prompt_request_pieces_with_scores(entries)
             return result
         except Exception as e:
             logger.exception(
@@ -161,7 +161,7 @@ def get_prompt_request_piece_by_memory_labels(
             conditions = [PromptMemoryEntry.labels.op("->>")(key) == value for key, value in memory_labels.items()]
             query_condition = and_(*conditions)
             entries = self.query_entries(PromptMemoryEntry, conditions=query_condition)
-            result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+            result = self.get_prompt_request_pieces_with_scores(entries)
             return result
         except Exception as e:
             logger.exception(
@@ -185,7 +185,7 @@ def _get_prompt_pieces_by_orchestrator(self, *, orchestrator_id: str) -> list[Pr
                 PromptMemoryEntry,
                 conditions=PromptMemoryEntry.orchestrator_identifier.op("->>")("id") == orchestrator_id,
             )  # type: ignore
-            result: list[PromptRequestPiece] = [entry.get_prompt_request_piece() for entry in entries]
+            result = self.get_prompt_request_pieces_with_scores(entries)
             return result
         except Exception as e:
             logger.exception(

pyrit/memory/duckdb_memory.py

@@ -3,15 +3,9 @@
 
 import csv
 import json
-from typing import Any, Dict, List, Union
-import uuid
-from datetime import datetime
 from pathlib import Path
-from collections.abc import MutableMapping
 
-from sqlalchemy.inspection import inspect
-
-from pyrit.memory.memory_models import Base
+from pyrit.models import PromptRequestPiece
 
 
 class MemoryExporter:
@@ -29,14 +23,13 @@ def __init__(self):
         }
 
     def export_data(
-        self, data: Union[List[Base], List[Dict]], *, file_path: Path = None, export_type: str = "json"
+        self, data: list[PromptRequestPiece], *, file_path: Path = None, export_type: str = "json"
     ):  # type: ignore
         """
         Exports the provided data to a file in the specified format.
 
         Args:
-            data (Union[List[Base], List[Dict]]): The data to be exported, typically a list of SQLAlchemy
-              model instances or as a list of dictionaries.
+            data (list[PromptRequestPiece]): The data to be exported, as a list of PromptRequestPiece instances.
             file_path (str): The full path, including the file name, where the data will be exported.
             export_type (str, Optional): The format for exporting data. Defaults to "json".
 
@@ -52,36 +45,37 @@ def export_data(
         else:
             raise ValueError(f"Unsupported export format: {export_type}")
 
-    def export_to_json(self, data: Union[List[Base], List[Dict]], file_path: Path = None) -> None:  # type: ignore
+    def export_to_json(self, data: list[PromptRequestPiece], file_path: Path = None) -> None:  # type: ignore
         """
         Exports the provided data to a JSON file at the specified file path.
         Each item in the data list, representing a row from the table,
         is converted to a dictionary before being written to the file.
 
         Args:
-            data (Union[List[Base], List[Dict]]): The data to be exported, as a list of SQLAlchemy model instances
-              or as a list of dictionaries.
+            data (list[PromptRequestPiece]): The data to be exported, as a list of PromptRequestPiece instances.
             file_path (Path): The full path, including the file name, where the data will be exported.
 
         Raises:
             ValueError: If no file_path is provided.
         """
         if not file_path:
             raise ValueError("Please provide a valid file path for exporting data.")
-
-        export_data = [self.model_to_dict(instance) if isinstance(instance, Base) else instance for instance in data]
+        if not data:
+            raise ValueError("No data to export.")
+        export_data = []
+        for piece in data:
+            export_data.append(piece.to_dict())
         with open(file_path, "w") as f:
             json.dump(export_data, f, indent=4)
 
-    def export_to_csv(self, data: Union[List[Base], List[Dict]], file_path: Path = None) -> None:  # type: ignore
+    def export_to_csv(self, data: list[PromptRequestPiece], file_path: Path = None) -> None:  # type: ignore
         """
         Exports the provided data to a CSV file at the specified file path.
         Each item in the data list, representing a row from the table,
         is converted to a dictionary before being written to the file.
 
         Args:
-            data (Union[List[Base], List[Dict]]): The data to be exported, as a list of SQLAlchemy model instances
-              or as a list of dictionaries.
+            data (list[PromptRequestPiece]): The data to be exported, as a list of PromptRequestPiece instances.
             file_path (Path): The full path, including the file name, where the data will be exported.
 
         Raises:
@@ -91,49 +85,11 @@ def export_to_csv(self, data: Union[List[Base], List[Dict]], file_path: Path = N
             raise ValueError("Please provide a valid file path for exporting data.")
         if not data:
             raise ValueError("No data to export.")
-
-        export_data = [
-            _flatten_dict(self.model_to_dict(instance)) if isinstance(instance, Base) else _flatten_dict(instance)
-            for instance in data
-        ]
+        export_data = []
+        for piece in data:
+            export_data.append(piece.to_dict())
         fieldnames = list(export_data[0].keys())
         with open(file_path, "w", newline="") as f:
             writer = csv.DictWriter(f, fieldnames=fieldnames)
             writer.writeheader()
             writer.writerows(export_data)
-
-    def model_to_dict(self, model_instance: Base):  # type: ignore
-        """
-        Converts an SQLAlchemy model instance into a dictionary, serializing
-        special data types such as UUID and datetime to string representations.
-        This ensures compatibility with JSON and other serialization formats.
-
-        Args:
-            model_instance: An instance of an SQLAlchemy model.
-
-        Returns:
-            A dictionary representation of the model instance, with special types serialized.
-        """
-        model_dict = {}
-        for column in inspect(model_instance.__class__).columns:
-            value = getattr(model_instance, column.name)
-            if isinstance(value, uuid.UUID):
-                # Convert UUID to string
-                model_dict[column.name] = str(value)
-            elif isinstance(value, datetime):
-                # Convert datetime to an ISO 8601 formatted string
-                model_dict[column.name] = value.isoformat()
-            else:
-                model_dict[column.name] = value
-        return model_dict
-
-
-def _flatten_dict(d: MutableMapping, parent_key: str = "", sep: str = ".") -> MutableMapping:
-    items: list[tuple[Any, Any]] = []
-    for k, v in d.items():
-        new_key = parent_key + sep + k if parent_key else k
-        if isinstance(v, MutableMapping):
-            items.extend(_flatten_dict(v, new_key, sep=sep).items())
-        else:
-            items.append((new_key, v))
-    return dict(items)