From b74fb381c1f16e79d27c89310c716a1b25754d9c Mon Sep 17 00:00:00 2001 From: biefan <70761325+biefan@users.noreply.github.com> Date: Tue, 17 Mar 2026 03:54:10 +0000 Subject: [PATCH 1/3] Respect export type in SQLite conversation exports --- pyrit/memory/sqlite_memory.py | 23 +++++- .../memory_interface/test_interface_export.py | 80 +++++++++++++++++++ 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 58ae9098ed..72f7818e6b 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -34,6 +34,14 @@ Model = TypeVar("Model") +class _ExportableConversationPiece: + def __init__(self, data: dict[str, Any]) -> None: + self._data = data + + def to_dict(self) -> dict[str, Any]: + return self._data + + class SQLiteMemory(MemoryInterface, metaclass=Singleton): """ A memory interface that uses SQLite as the backend database. @@ -418,9 +426,18 @@ def export_conversations( piece_data["scores"] = [score.to_dict() for score in piece_scores] merged_data.append(piece_data) - # Export to JSON manually since the exporter expects objects but we have dicts - with open(file_path, "w") as f: - json.dump(merged_data, f, indent=4) + if not merged_data: + if export_type == "json": + with open(file_path, "w", encoding="utf-8") as f: + json.dump(merged_data, f, indent=4) + elif export_type in self.exporter.export_strategies: + file_path.write_text("", encoding="utf-8") + else: + raise ValueError(f"Unsupported export format: {export_type}") + return file_path + + exportable_pieces = [_ExportableConversationPiece(data=piece_data) for piece_data in merged_data] + self.exporter.export_data(exportable_pieces, file_path=file_path, export_type=export_type) return file_path def print_schema(self) -> None: diff --git a/tests/unit/memory/memory_interface/test_interface_export.py b/tests/unit/memory/memory_interface/test_interface_export.py index 14064b718d..aaf4dc57b3 100644 --- a/tests/unit/memory/memory_interface/test_interface_export.py +++ b/tests/unit/memory/memory_interface/test_interface_export.py @@ -3,6 +3,7 @@ import os import tempfile +import csv from collections.abc import Sequence from pathlib import Path from unittest.mock import MagicMock, patch @@ -151,3 +152,82 @@ def test_export_all_conversations_with_scores_empty_data(sqlite_instance: Memory # Clean up the temp file if file_path.exists(): os.remove(file_path) + + +def test_export_all_conversations_with_scores_csv_format(sqlite_instance: MemoryInterface): + sqlite_instance.exporter = MemoryExporter() + + with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: + file_path = Path(temp_file.name) + temp_file.close() + + try: + with ( + patch.object(sqlite_instance, "get_message_pieces") as mock_get_pieces, + patch.object(sqlite_instance, "get_prompt_scores") as mock_get_scores, + ): + mock_piece = MagicMock() + mock_piece.id = "piece_id_1234" + mock_piece.to_dict.return_value = { + "id": "piece_id_1234", + "converted_value": "sample piece", + } + + mock_score = MagicMock() + mock_score.message_piece_id = "piece_id_1234" + mock_score.to_dict.return_value = {"message_piece_id": "piece_id_1234", "score_value": 10} + + mock_get_pieces.return_value = [mock_piece] + mock_get_scores.return_value = [mock_score] + + sqlite_instance.export_conversations(file_path=file_path, export_type="csv") + + with open(file_path, newline="") as exported_file: + reader = csv.DictReader(exported_file) + assert reader.fieldnames == ["id", "converted_value", "scores"] + rows = list(reader) + + assert len(rows) == 1 + assert rows[0]["id"] == "piece_id_1234" + assert rows[0]["converted_value"] == "sample piece" + finally: + if file_path.exists(): + os.remove(file_path) + + +def test_export_all_conversations_with_scores_markdown_format(sqlite_instance: MemoryInterface): + sqlite_instance.exporter = MemoryExporter() + + with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as temp_file: + file_path = Path(temp_file.name) + temp_file.close() + + try: + with ( + patch.object(sqlite_instance, "get_message_pieces") as mock_get_pieces, + patch.object(sqlite_instance, "get_prompt_scores") as mock_get_scores, + ): + mock_piece = MagicMock() + mock_piece.id = "piece_id_1234" + mock_piece.to_dict.return_value = { + "id": "piece_id_1234", + "converted_value": "sample piece", + } + + mock_score = MagicMock() + mock_score.message_piece_id = "piece_id_1234" + mock_score.to_dict.return_value = {"message_piece_id": "piece_id_1234", "score_value": 10} + + mock_get_pieces.return_value = [mock_piece] + mock_get_scores.return_value = [mock_score] + + sqlite_instance.export_conversations(file_path=file_path, export_type="md") + + exported_content = file_path.read_text(encoding="utf-8") + + assert exported_content.startswith("| id | converted_value | scores |") + assert "piece_id_1234" in exported_content + assert "sample piece" in exported_content + finally: + if file_path.exists(): + os.remove(file_path) From 23cbfcec6f45f45098b1bde55bd42808a67fc65d Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 10 Apr 2026 19:40:12 -0700 Subject: [PATCH 2/3] Parametrize export format tests per review feedback Combine separate CSV and Markdown test functions into a single parametrized test covering json, csv, and md export types. Move inline imports to file top. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../memory_interface/test_interface_export.py | 76 ++++++------------- 1 file changed, 25 insertions(+), 51 deletions(-) diff --git a/tests/unit/memory/memory_interface/test_interface_export.py b/tests/unit/memory/memory_interface/test_interface_export.py index aaf4dc57b3..42d4e6d80f 100644 --- a/tests/unit/memory/memory_interface/test_interface_export.py +++ b/tests/unit/memory/memory_interface/test_interface_export.py @@ -1,13 +1,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import csv +import json import os import tempfile -import csv from collections.abc import Sequence from pathlib import Path from unittest.mock import MagicMock, patch +import pytest + from pyrit.common.path import DB_DATA_PATH from pyrit.memory import MemoryExporter, MemoryInterface from pyrit.models import MessagePiece @@ -104,8 +107,6 @@ def test_export_all_conversations_with_scores_correct_data(sqlite_instance: Memo assert file_path.exists() # Read and verify the exported JSON content - import json - with open(file_path) as f: exported_data = json.load(f) @@ -142,8 +143,6 @@ def test_export_all_conversations_with_scores_empty_data(sqlite_instance: Memory assert file_path.exists() # Read and verify the exported JSON content is empty - import json - with open(file_path) as f: exported_data = json.load(f) @@ -154,51 +153,13 @@ def test_export_all_conversations_with_scores_empty_data(sqlite_instance: Memory os.remove(file_path) -def test_export_all_conversations_with_scores_csv_format(sqlite_instance: MemoryInterface): - sqlite_instance.exporter = MemoryExporter() - - with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: - file_path = Path(temp_file.name) - temp_file.close() - - try: - with ( - patch.object(sqlite_instance, "get_message_pieces") as mock_get_pieces, - patch.object(sqlite_instance, "get_prompt_scores") as mock_get_scores, - ): - mock_piece = MagicMock() - mock_piece.id = "piece_id_1234" - mock_piece.to_dict.return_value = { - "id": "piece_id_1234", - "converted_value": "sample piece", - } - - mock_score = MagicMock() - mock_score.message_piece_id = "piece_id_1234" - mock_score.to_dict.return_value = {"message_piece_id": "piece_id_1234", "score_value": 10} - - mock_get_pieces.return_value = [mock_piece] - mock_get_scores.return_value = [mock_score] - - sqlite_instance.export_conversations(file_path=file_path, export_type="csv") - - with open(file_path, newline="") as exported_file: - reader = csv.DictReader(exported_file) - assert reader.fieldnames == ["id", "converted_value", "scores"] - rows = list(reader) - - assert len(rows) == 1 - assert rows[0]["id"] == "piece_id_1234" - assert rows[0]["converted_value"] == "sample piece" - finally: - if file_path.exists(): - os.remove(file_path) - - -def test_export_all_conversations_with_scores_markdown_format(sqlite_instance: MemoryInterface): +@pytest.mark.parametrize("export_type, suffix", [("json", ".json"), ("csv", ".csv"), ("md", ".md")]) +def test_export_all_conversations_with_scores_respects_export_type( + sqlite_instance: MemoryInterface, export_type: str, suffix: str +): sqlite_instance.exporter = MemoryExporter() - with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as temp_file: + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: file_path = Path(temp_file.name) temp_file.close() @@ -221,13 +182,26 @@ def test_export_all_conversations_with_scores_markdown_format(sqlite_instance: M mock_get_pieces.return_value = [mock_piece] mock_get_scores.return_value = [mock_score] - sqlite_instance.export_conversations(file_path=file_path, export_type="md") + sqlite_instance.export_conversations(file_path=file_path, export_type=export_type) + assert file_path.exists() exported_content = file_path.read_text(encoding="utf-8") - - assert exported_content.startswith("| id | converted_value | scores |") assert "piece_id_1234" in exported_content assert "sample piece" in exported_content + + if export_type == "json": + exported_data = json.loads(exported_content) + assert len(exported_data) == 1 + assert exported_data[0]["id"] == "piece_id_1234" + elif export_type == "csv": + with open(file_path, newline="") as exported_file: + reader = csv.DictReader(exported_file) + assert reader.fieldnames == ["id", "converted_value", "scores"] + rows = list(reader) + assert len(rows) == 1 + assert rows[0]["id"] == "piece_id_1234" + elif export_type == "md": + assert exported_content.startswith("| id | converted_value | scores |") finally: if file_path.exists(): os.remove(file_path) From e35d6faf7b4044f247699f37c46a2086164792a3 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 10 Apr 2026 19:52:41 -0700 Subject: [PATCH 3/3] Fix ruff and mypy pre-commit errors Add ValueError to export_conversations docstring (DOC501). Use cast for _ExportableConversationPiece list to satisfy mypy arg-type check. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/memory/sqlite_memory.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index ba711b9fe9..9dd3561f4a 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -8,7 +8,7 @@ from contextlib import closing, suppress from datetime import datetime from pathlib import Path -from typing import Any, Optional, TypeVar, Union +from typing import Any, Optional, TypeVar, Union, cast from sqlalchemy import and_, create_engine, func, or_, text from sqlalchemy.engine.base import Engine @@ -482,6 +482,9 @@ def export_conversations( Returns: Path: The path to the exported file. + + Raises: + ValueError: If the specified export format is not supported. """ # Import here to avoid circular import issues from pyrit.memory.memory_exporter import MemoryExporter @@ -541,7 +544,9 @@ def export_conversations( return file_path exportable_pieces = [_ExportableConversationPiece(data=piece_data) for piece_data in merged_data] - self.exporter.export_data(exportable_pieces, file_path=file_path, export_type=export_type) + self.exporter.export_data( + cast("list[MessagePiece]", exportable_pieces), file_path=file_path, export_type=export_type + ) return file_path def print_schema(self) -> None: