FEAT: Extend data exporter to support Markdown (Azure#1033) (Azure#1042)

Co-authored-by: 1twodrei <[email protected]>
Co-authored-by: rlundeen2 <[email protected]>
Co-authored-by: Richard Lundeen <[email protected]>

Author: 3 people

pyrit/memory/memory_exporter.py

@@ -20,6 +20,7 @@ def __init__(self):
         self.export_strategies = {
             "json": self.export_to_json,
             "csv": self.export_to_csv,
+            "md": self.export_to_markdown,
             # Future formats can be added here
         }
 
@@ -94,3 +95,27 @@ def export_to_csv(self, data: list[PromptRequestPiece], file_path: Path = None)
             writer = csv.DictWriter(f, fieldnames=fieldnames)
             writer.writeheader()
             writer.writerows(export_data)
+
+    def export_to_markdown(self, data: list[PromptRequestPiece], file_path: Path = None) -> None:  # type: ignore
+        """
+        Exports the provided data to a Markdown file at the specified file path.
+        Each item in the data list is converted to a dictionary and formatted as a table.
+
+        Args:
+            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 or if there is no data to export.
+        """
+        if not file_path:
+            raise ValueError("Please provide a valid file path for exporting data.")
+        if not data:
+            raise ValueError("No data to export.")
+        export_data = [piece.to_dict() for piece in data]
+        fieldnames = list(export_data[0].keys())
+        with open(file_path, "w", newline="") as f:
+            f.write(f"| {' | '.join(fieldnames)} |\n")
+            f.write(f"| {' | '.join(['---'] * len(fieldnames))} |\n")
+            for row in export_data:
+                f.write(f"| {' | '.join(str(row[field]) for field in fieldnames)} |\n")

tests/unit/memory/test_memory_exporter.py

@@ -90,3 +90,13 @@ def test_export_to_json_data_with_conversations(tmp_path, export_type):
     assert content[1]["role"] == "assistant"
     assert content[1]["converted_value"] == "I'm fine, thank you!"
     assert content[1]["conversation_id"] == conversation_id
+
+
+@pytest.mark.parametrize("export_type", ["json", "csv", "md"])
+def test_export_data_creates_file(tmp_path, export_type):
+    exporter = MemoryExporter()
+    file_path = tmp_path / f"conversations.{export_type}"
+    sample_conversation_entries = get_sample_conversations()
+    exporter.export_data(data=sample_conversation_entries, file_path=file_path, export_type=export_type)
+
+    assert file_path.exists()