diff --git a/pyrit/backend/mappers/__init__.py b/pyrit/backend/mappers/__init__.py index 55c8b2af83..310b04e916 100644 --- a/pyrit/backend/mappers/__init__.py +++ b/pyrit/backend/mappers/__init__.py @@ -8,6 +8,7 @@ Centralizes all translation logic so domain models can evolve independently of the API contract. """ +from pyrit.backend.mappers._preview import format_last_message_preview from pyrit.backend.mappers.attack_mappers import ( attack_result_to_summary, pyrit_messages_to_dto_async, @@ -25,6 +26,7 @@ __all__ = [ "attack_result_to_summary", "converter_object_to_instance", + "format_last_message_preview", "pyrit_messages_to_dto_async", "pyrit_scores_to_dto", "request_piece_to_pyrit_message_piece", diff --git a/pyrit/backend/mappers/_preview.py b/pyrit/backend/mappers/_preview.py new file mode 100644 index 0000000000..33e5f1ea73 --- /dev/null +++ b/pyrit/backend/mappers/_preview.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Presentation-layer formatter for ``ConversationStats.last_message_preview``. + +Lives in the backend mapper package because the formatting it produces +(``[Image: ]`` etc.) is purely a display concern for the GUI API +responses — the memory layer stays data-agnostic and just stores the raw +value + data type. + +The motivating bug: ``converted_value`` for media-path data types +(``image_path`` / ``audio_path`` / ``video_path`` / ``binary_path``) is a +filesystem path or blob URL. Rendering it raw in the Attack History preview +leaks the absolute on-disk location of memory artifacts +(e.g. ``C:\\Users\\\\git\\PyRIT\\dbdata\\...\\1780.mp3``). +""" + +from pathlib import PureWindowsPath +from urllib.parse import urlparse + +from pyrit.models import MEDIA_PATH_DATA_TYPES, ConversationStats + +# Friendly label per media-path data type. Kept here next to the formatter +# so adding a new media type only requires updating one place. +_MEDIA_LABEL: dict[str, str] = { + "image_path": "Image", + "audio_path": "Audio", + "video_path": "Video", + "binary_path": "File", +} + + +def _derive_basename(value: str) -> str | None: + """ + Return a display-safe basename for *value*. + + Args: + value: A filesystem path, URL, or other reference. + + Returns: + The basename (filename portion) of *value*, or ``None`` if one can't + be derived (e.g. data URI, empty value). + """ + if not value or value.startswith("data:"): + return None + if value.startswith(("http://", "https://")): + # Strip query string (e.g. SAS tokens) before taking the basename. + parsed = urlparse(value) + name = PureWindowsPath(parsed.path).name + return name or None + # Local path — PureWindowsPath treats both ``/`` and ``\`` as separators, + # so Windows-style paths stored from a Windows host are split correctly + # even when this code runs on a POSIX host (CI, Linux deployments). + return PureWindowsPath(value).name or None + + +def format_last_message_preview( + *, + value: str | None, + data_type: str | None, + max_len: int = ConversationStats.PREVIEW_MAX_LEN, +) -> str | None: + """ + Build a display string for ``ConversationStats.last_message_preview``. + + Media-path data types are rendered as ``[Image: ]`` (and + variants) so the absolute filesystem path of memory artifacts is never + exposed through API responses or UI previews. Text-like data types pass + through with truncation and an ellipsis suffix when they exceed + *max_len*. + + Args: + value: Raw ``converted_value`` for the last piece (or ``None``). + data_type: ``converted_value_data_type`` for that piece. ``None`` + falls back to the text path. + max_len: Maximum length for text previews before truncation. + + Returns: + The formatted preview string, or ``None`` when there is nothing + meaningful to show. + """ + if data_type in MEDIA_PATH_DATA_TYPES: + # MEDIA_PATH_DATA_TYPES guarantees ``data_type`` is a key in + # ``_MEDIA_LABEL`` — both are derived from the same source list. + label = _MEDIA_LABEL[data_type] + basename = _derive_basename(value or "") + return f"[{label}: {basename}]" if basename else f"[{label}]" + + if not value: + return None + + if len(value) > max_len: + return value[:max_len] + "..." + return value diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index 18282fe22b..5807c27bef 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -24,6 +24,7 @@ from azure.storage.blob import ContainerSasPermissions, generate_container_sas from azure.storage.blob.aio import BlobServiceClient +from pyrit.backend.mappers._preview import format_last_message_preview from pyrit.backend.models.attacks import ( AddMessageRequest, AttackSummary, @@ -35,7 +36,7 @@ TargetInfo, ) from pyrit.common.deprecation import print_deprecation_message -from pyrit.models import AttackResult, ChatMessageRole, PromptDataType +from pyrit.models import MEDIA_PATH_DATA_TYPES, AttackResult, ChatMessageRole, PromptDataType from pyrit.models import Message as PyritMessage from pyrit.models import MessagePiece as PyritMessagePiece from pyrit.models import Score as PyritScore @@ -50,9 +51,6 @@ # Domain → DTO (for API responses) # ============================================================================ -# Media data types whose values are file paths (local or Azure Blob URLs) -_MEDIA_PATH_TYPES = frozenset({"image_path", "audio_path", "video_path", "binary_path"}) - # --------------------------------------------------------------------------- # Azure Blob SAS token cache # --------------------------------------------------------------------------- @@ -172,7 +170,7 @@ def _resolve_media_url(*, value: Optional[str], data_type: str) -> Optional[str] The value unchanged for non-media types, a ``/api/media?path=...`` URL for local file paths, or the original value for blob URLs / data URIs. """ - if not value or data_type not in _MEDIA_PATH_TYPES: + if not value or data_type not in MEDIA_PATH_DATA_TYPES: return value # Already a URL or data URI — pass through if value.startswith(("http://", "https://", "data:")): @@ -227,7 +225,10 @@ def attack_result_to_summary( AttackSummary DTO ready for the API response. """ message_count = stats.message_count - last_preview = stats.last_message_preview + last_preview = format_last_message_preview( + value=stats.last_message_preview, + data_type=stats.last_message_data_type, + ) # Merge attack-result labels with conversation-level labels. # Conversation labels take precedence on key collision. @@ -297,7 +298,9 @@ def pyrit_scores_to_dto(scores: list[PyritScore]) -> list[Score]: return [ Score( score_id=str(score.id), - scorer_type=score.scorer_class_identifier.class_name, + scorer_type=( + score.scorer_class_identifier.class_name or "Unknown" if score.scorer_class_identifier else "Unknown" + ), score_type=score.score_type, score_value=score.score_value, score_category=score.score_category, diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index c5502bb018..74f96ff9f3 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -24,8 +24,9 @@ from typing import Any, Literal, cast from urllib.parse import parse_qs, urlparse -from pyrit.backend.mappers.attack_mappers import ( +from pyrit.backend.mappers import ( attack_result_to_summary, + format_last_message_preview, pyrit_messages_to_dto_async, request_piece_to_pyrit_message_piece, request_to_pyrit_message, @@ -177,11 +178,13 @@ async def list_attacks_async( total_count = (main_stats.message_count if main_stats else 0) + sum(s.message_count for s in pruned_stats) preview = main_stats.last_message_preview if main_stats else None + preview_data_type = main_stats.last_message_data_type if main_stats else None conv_labels = (main_stats.labels if main_stats else None) or {} merged = ConversationStats( message_count=total_count, last_message_preview=preview, + last_message_data_type=preview_data_type, labels=conv_labels, ) @@ -419,7 +422,10 @@ async def get_conversations_async(self, *, attack_result_id: str) -> AttackConve ConversationSummary( conversation_id=conv_id, message_count=stats.message_count if stats else 0, - last_message_preview=stats.last_message_preview if stats else None, + last_message_preview=format_last_message_preview( + value=stats.last_message_preview if stats else None, + data_type=stats.last_message_data_type if stats else None, + ), created_at=created_at, ) ) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index a04ab0bbee..6723ae2842 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -615,18 +615,23 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str placeholders = ", ".join(f":cid{i}" for i in range(len(conversation_ids))) params = {f"cid{i}": cid for i, cid in enumerate(conversation_ids)} - max_len = ConversationStats.PREVIEW_MAX_LEN sql = text( f""" SELECT pme.conversation_id, COUNT(DISTINCT pme.sequence) AS msg_count, ( - SELECT TOP 1 LEFT(p2.converted_value, {max_len + 3}) + SELECT TOP 1 LEFT(p2.converted_value, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) FROM "PromptMemoryEntries" p2 WHERE p2.conversation_id = pme.conversation_id ORDER BY p2.sequence DESC, p2.id DESC ) AS last_preview, + ( + SELECT TOP 1 p2b.converted_value_data_type + FROM "PromptMemoryEntries" p2b + WHERE p2b.conversation_id = pme.conversation_id + ORDER BY p2b.sequence DESC, p2b.id DESC + ) AS last_data_type, ( SELECT TOP 1 p3.labels FROM "PromptMemoryEntries" p3 @@ -648,11 +653,7 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str result: dict[str, ConversationStats] = {} for row in rows: - conv_id, msg_count, last_preview, raw_labels, raw_created_at = row - - preview = None - if last_preview: - preview = last_preview[:max_len] + "..." if len(last_preview) > max_len else last_preview + conv_id, msg_count, last_preview, last_data_type, raw_labels, raw_created_at = row labels: dict[str, str] = {} if raw_labels and raw_labels not in ("null", "{}"): @@ -668,7 +669,8 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str result[conv_id] = ConversationStats( message_count=msg_count, - last_message_preview=preview, + last_message_preview=last_preview, + last_message_data_type=last_data_type, labels=labels, created_at=created_at, ) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index d39897046a..461d2b871b 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -739,19 +739,25 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str placeholders = ", ".join(f":cid{i}" for i in range(len(conversation_ids))) params = {f"cid{i}": cid for i, cid in enumerate(conversation_ids)} - max_len = ConversationStats.PREVIEW_MAX_LEN sql = text( f""" SELECT pme.conversation_id, COUNT(DISTINCT pme.sequence) AS msg_count, ( - SELECT SUBSTR(p2.converted_value, 1, {max_len + 3}) + SELECT SUBSTR(p2.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) FROM "PromptMemoryEntries" p2 WHERE p2.conversation_id = pme.conversation_id ORDER BY p2.sequence DESC, p2.id DESC LIMIT 1 ) AS last_preview, + ( + SELECT p2b.converted_value_data_type + FROM "PromptMemoryEntries" p2b + WHERE p2b.conversation_id = pme.conversation_id + ORDER BY p2b.sequence DESC, p2b.id DESC + LIMIT 1 + ) AS last_data_type, ( SELECT p3.labels FROM "PromptMemoryEntries" p3 @@ -774,11 +780,7 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str result: dict[str, ConversationStats] = {} for row in rows: - conv_id, msg_count, last_preview, raw_labels, raw_created_at = row - - preview = None - if last_preview: - preview = last_preview[:max_len] + "..." if len(last_preview) > max_len else last_preview + conv_id, msg_count, last_preview, last_data_type, raw_labels, raw_created_at = row labels: dict[str, str] = {} if raw_labels and raw_labels not in ("null", "{}"): @@ -794,7 +796,8 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str result[conv_id] = ConversationStats( message_count=msg_count, - last_message_preview=preview, + last_message_preview=last_preview, + last_message_data_type=last_data_type, labels=labels, created_at=created_at, ) diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index cb071a0b9a..55e7110181 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -62,7 +62,14 @@ snake_case_to_class_name, validate_registry_name, ) -from pyrit.models.literals import ChatMessageRole, Modality, PromptDataType, PromptResponseError, SeedType +from pyrit.models.literals import ( + MEDIA_PATH_DATA_TYPES, + ChatMessageRole, + Modality, + PromptDataType, + PromptResponseError, + SeedType, +) from pyrit.models.messages import ( Message, MessagePiece, @@ -141,6 +148,7 @@ "IdentifierFilter", "IdentifierType", "ImagePathDataTypeSerializer", + "MEDIA_PATH_DATA_TYPES", "Message", "MessagePiece", "Modality", diff --git a/pyrit/models/conversation_stats.py b/pyrit/models/conversation_stats.py index 22dcefd6ea..67b09e24be 100644 --- a/pyrit/models/conversation_stats.py +++ b/pyrit/models/conversation_stats.py @@ -6,6 +6,8 @@ from pydantic import BaseModel, ConfigDict, Field +from pyrit.models.literals import PromptDataType + class ConversationStats(BaseModel): """ @@ -17,8 +19,17 @@ class ConversationStats(BaseModel): model_config = ConfigDict(frozen=True) PREVIEW_MAX_LEN: ClassVar[int] = 100 + PREVIEW_FETCH_MAX_LEN: ClassVar[int] = 1024 + """ + Upper bound (in characters) for the raw ``last_message_preview`` value + fetched from storage. Larger than ``PREVIEW_MAX_LEN`` so that downstream + presentation code (see ``pyrit.backend.mappers._preview``) has enough + characters to extract a basename from a long media path or signed blob + URL before applying display-level truncation. + """ message_count: int = 0 last_message_preview: Optional[str] = None + last_message_data_type: Optional[PromptDataType] = None labels: dict[str, str] = Field(default_factory=dict) created_at: Optional[datetime] = None diff --git a/pyrit/models/literals.py b/pyrit/models/literals.py index a61b65a09c..4d86b7d069 100644 --- a/pyrit/models/literals.py +++ b/pyrit/models/literals.py @@ -18,6 +18,12 @@ "function_call_output", ] +# Subset of ``PromptDataType`` values whose stored ``value`` is a path or URL +# pointing at media content (rather than the content itself). Useful for +# treating these specially — e.g. avoiding raw filesystem-path leaks in API +# previews, or signing blob storage URLs before exposing them to the frontend. +MEDIA_PATH_DATA_TYPES: frozenset[PromptDataType] = frozenset({"image_path", "audio_path", "video_path", "binary_path"}) + """ The type of the error in the prompt response blocked: blocked by an external filter e.g. Azure Filters diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index bf63d17af5..dbc12ab70b 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -415,6 +415,26 @@ async def test_list_attacks_includes_labels_in_summary(self, attack_service, moc assert len(result.items) == 1 assert result.items[0].labels == {"env": "prod", "team": "red", "test_ar_label": "test_ar_value"} + async def test_list_attacks_formats_media_preview(self, attack_service, mock_memory) -> None: + """list_attacks AttackSummary previews must not leak absolute media paths.""" + ar = make_attack_result(conversation_id="attack-1") + mock_memory.get_attack_results.return_value = [ar] + path = r"C:\Users\someone\PyRIT\dbdata\prompt-memory-entries\images\1780010098266691.png" + mock_memory.get_conversation_stats.return_value = { + "attack-1": ConversationStats( + message_count=1, + last_message_preview=path, + last_message_data_type="image_path", + ), + } + + result = await attack_service.list_attacks_async() + + assert len(result.items) == 1 + preview = result.items[0].last_message_preview + assert preview == "[Image: 1780010098266691.png]" + assert "C:\\" not in (preview or "") + async def test_list_attacks_filters_by_labels_directly(self, attack_service, mock_memory) -> None: """Test that label filters are passed directly to the DB query (no legacy expansion).""" ar = make_attack_result(conversation_id="attack-canonical") @@ -1709,6 +1729,26 @@ async def test_returns_main_conversation_only(self, attack_service, mock_memory) assert len(result.conversations) == 1 assert result.conversations[0].message_count == 2 + async def test_conversation_summary_formats_media_preview(self, attack_service, mock_memory): + """ConversationSummary previews must not leak absolute media paths.""" + ar = make_attack_result(conversation_id="attack-1") + mock_memory.get_attack_results.return_value = [ar] + path = r"C:\Users\someone\PyRIT\dbdata\prompt-memory-entries\audio\1780010098266691.mp3" + mock_memory.get_conversation_stats.return_value = { + "attack-1": ConversationStats( + message_count=1, + last_message_preview=path, + last_message_data_type="audio_path", + ), + } + + result = await attack_service.get_conversations_async(attack_result_id="attack-1") + + assert result is not None + preview = result.conversations[0].last_message_preview + assert preview == "[Audio: 1780010098266691.mp3]" + assert "C:\\" not in (preview or "") + async def test_returns_main_and_related_conversations(self, attack_service, mock_memory): """Should return main and PRUNED conversations sorted by timestamp.""" from pyrit.models.conversation_reference import ConversationReference, ConversationType diff --git a/tests/unit/backend/test_mappers.py b/tests/unit/backend/test_mappers.py index b38c7fa6e5..2d6e4da37c 100644 --- a/tests/unit/backend/test_mappers.py +++ b/tests/unit/backend/test_mappers.py @@ -154,11 +154,11 @@ def test_empty_pieces_gives_zero_messages(self) -> None: assert summary.message_count == 0 assert summary.last_message_preview is None - def test_last_message_preview_truncated(self) -> None: - """Test that long messages are truncated in stats.""" + def test_last_message_preview_truncates_long_raw_text(self) -> None: + """The mapper applies the preview formatter, which truncates long raw text.""" ar = _make_attack_result() long_text = "x" * 200 - stats = ConversationStats(message_count=1, last_message_preview=long_text[:100] + "...") + stats = ConversationStats(message_count=1, last_message_preview=long_text, last_message_data_type="text") summary = attack_result_to_summary(ar, stats=stats) @@ -166,6 +166,31 @@ def test_last_message_preview_truncated(self) -> None: assert len(summary.last_message_preview) == 103 # 100 + "..." assert summary.last_message_preview.endswith("...") + @pytest.mark.parametrize( + ("data_type", "expected"), + [ + ("image_path", "[Image: 1780010098266691.png]"), + ("audio_path", "[Audio: 1780010098266691.png]"), + ("video_path", "[Video: 1780010098266691.png]"), + ("binary_path", "[File: 1780010098266691.png]"), + ], + ) + def test_media_last_message_preview_hides_absolute_path(self, data_type: str, expected: str) -> None: + """The mapper renders media-type previews as friendly labels rather + than leaking the raw on-disk path it receives from memory.""" + ar = _make_attack_result() + path = r"C:\Users\someone\git\PyRIT\dbdata\prompt-memory-entries\media\1780010098266691.png" + stats = ConversationStats( + message_count=1, + last_message_preview=path, + last_message_data_type=data_type, + ) + + summary = attack_result_to_summary(ar, stats=stats) + + assert summary.last_message_preview == expected + assert "C:\\" not in (summary.last_message_preview or "") + def test_labels_are_mapped(self) -> None: """Test that labels are derived from stats.""" ar = _make_attack_result() diff --git a/tests/unit/backend/test_preview.py b/tests/unit/backend/test_preview.py new file mode 100644 index 0000000000..29e1dc498b --- /dev/null +++ b/tests/unit/backend/test_preview.py @@ -0,0 +1,105 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit tests for ``pyrit.backend.mappers.format_last_message_preview``.""" + +import pytest + +from pyrit.backend.mappers import format_last_message_preview +from pyrit.models import ConversationStats + + +class TestFormatLastMessagePreview: + def test_text_short_value_passes_through(self) -> None: + result = format_last_message_preview(value="hello", data_type="text", max_len=100) + assert result == "hello" + + def test_text_long_value_is_truncated_with_ellipsis(self) -> None: + long_text = "x" * 200 + result = format_last_message_preview(value=long_text, data_type="text", max_len=100) + assert result is not None + assert len(result) == 103 + assert result.endswith("...") + assert result.startswith("x" * 100) + + def test_text_none_value_returns_none(self) -> None: + assert format_last_message_preview(value=None, data_type="text", max_len=100) is None + + def test_text_empty_value_returns_none(self) -> None: + assert format_last_message_preview(value="", data_type="text", max_len=100) is None + + def test_unknown_data_type_treated_as_text(self) -> None: + result = format_last_message_preview(value="hello", data_type=None, max_len=100) + assert result == "hello" + + def test_default_max_len_matches_conversation_stats_contract(self) -> None: + # The formatter's default truncation length should track the model + # constant so callers don't have to plumb it through manually. + long_text = "y" * (ConversationStats.PREVIEW_MAX_LEN + 50) + result = format_last_message_preview(value=long_text, data_type="text") + assert result is not None + assert len(result) == ConversationStats.PREVIEW_MAX_LEN + 3 + assert result.endswith("...") + + @pytest.mark.parametrize( + ("data_type", "label"), + [ + ("image_path", "Image"), + ("audio_path", "Audio"), + ("video_path", "Video"), + ("binary_path", "File"), + ], + ) + def test_media_windows_absolute_path_renders_basename_only(self, data_type: str, label: str) -> None: + path = r"C:\Users\someone\git\PyRIT\dbdata\prompt-memory-entries\audio\1780010098266691.mp3" + result = format_last_message_preview(value=path, data_type=data_type, max_len=100) + assert result == f"[{label}: 1780010098266691.mp3]" + assert "C:\\" not in (result or "") + assert "Users" not in (result or "") + + def test_media_posix_absolute_path_renders_basename_only(self) -> None: + path = "/home/someone/PyRIT/dbdata/prompt-memory-entries/images/abcdef.png" + result = format_last_message_preview(value=path, data_type="image_path", max_len=100) + assert result == "[Image: abcdef.png]" + assert "/home/" not in (result or "") + + def test_media_relative_path_renders_basename(self) -> None: + result = format_last_message_preview(value="audio/foo.mp3", data_type="audio_path", max_len=100) + assert result == "[Audio: foo.mp3]" + + def test_media_azure_blob_url_strips_query_and_keeps_filename(self) -> None: + url = "https://acct.blob.core.windows.net/container/folder/file.png?sv=2024-01-01&sig=secrettoken" + result = format_last_message_preview(value=url, data_type="image_path", max_len=100) + assert result == "[Image: file.png]" + assert "sig=" not in (result or "") + assert "blob.core.windows.net" not in (result or "") + + def test_media_empty_value_falls_back_to_label_only(self) -> None: + result = format_last_message_preview(value="", data_type="image_path", max_len=100) + assert result == "[Image]" + + def test_media_none_value_falls_back_to_label_only(self) -> None: + result = format_last_message_preview(value=None, data_type="audio_path", max_len=100) + assert result == "[Audio]" + + def test_media_data_uri_falls_back_to_label_only(self) -> None: + # Defensive: data URIs aren't expected as stored converted_value for + # media-path types, but if one shows up we should not try to derive a + # nonsensical basename from it. + result = format_last_message_preview( + value="data:image/png;base64,iVBORw0KGgo=", data_type="image_path", max_len=100 + ) + assert result == "[Image]" + + def test_media_long_path_basename_not_truncated(self) -> None: + # Even with a 100-char text limit, the basename label should not be + # truncated. Memory layer fetches up to PREVIEW_FETCH_MAX_LEN chars so + # the basename survives even very deep paths. + deep = "C:\\very\\deep\\nested\\directory\\structure\\that\\is\\quite\\long\\file_name_that_is_also_long.png" + result = format_last_message_preview(value=deep, data_type="image_path", max_len=20) + assert result == "[Image: file_name_that_is_also_long.png]" + + def test_preview_fetch_max_len_contract_is_generous(self) -> None: + # Sanity check on the model-side constant: must be large enough to fit + # realistic filesystem paths and signed blob URLs in a single fetch. + assert ConversationStats.PREVIEW_FETCH_MAX_LEN >= 512 diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index 9b5bacd33d..3eddf11133 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -793,19 +793,24 @@ def test_get_conversation_stats_returns_labels(sqlite_instance): assert result[conv_id].labels == {"env": "prod", "source": "gui"} -def test_get_conversation_stats_preview_truncates(sqlite_instance): - """Test that last_message_preview is truncated to 100 chars + ellipsis.""" +def test_get_conversation_stats_preview_caps_raw_value_at_fetch_limit(sqlite_instance): + """Memory caps the raw last_message_preview at PREVIEW_FETCH_MAX_LEN. + + Display-level truncation to PREVIEW_MAX_LEN happens later in the backend + mapper. This test verifies the storage-fetch contract: very long values + are bounded so a multi-MB text response doesn't bloat ``ConversationStats``. + """ import uuid - from pyrit.models import MessagePiece + from pyrit.models import ConversationStats, MessagePiece conv_id = str(uuid.uuid4()) - long_text = "x" * 200 + huge_text = "x" * (ConversationStats.PREVIEW_FETCH_MAX_LEN * 3) piece = MessagePiece( role="assistant", - original_value=long_text, + original_value=huge_text, original_value_data_type="text", - converted_value=long_text, + converted_value=huge_text, converted_value_data_type="text", conversation_id=conv_id, sequence=0, @@ -817,8 +822,9 @@ def test_get_conversation_stats_preview_truncates(sqlite_instance): assert conv_id in result preview = result[conv_id].last_message_preview assert preview is not None - assert len(preview) == 103 # 100 chars + "..." - assert preview.endswith("...") + assert len(preview) == ConversationStats.PREVIEW_FETCH_MAX_LEN + assert preview == "x" * ConversationStats.PREVIEW_FETCH_MAX_LEN + assert result[conv_id].last_message_data_type == "text" def test_get_conversation_stats_batches_multiple_conversations(sqlite_instance): @@ -852,6 +858,76 @@ def test_get_conversation_stats_batches_multiple_conversations(sqlite_instance): assert result[conv_ids[2]].message_count == 3 +@pytest.mark.parametrize( + "data_type", + ["image_path", "audio_path", "video_path", "binary_path"], +) +def test_get_conversation_stats_returns_media_data_type(sqlite_instance, data_type): + """Memory exposes the raw value + data type for the last piece — the + backend mapper handles display formatting. Verifies the data type is + propagated so downstream consumers can render media previews safely.""" + import uuid + + from pyrit.models import MessagePiece + + conv_id = str(uuid.uuid4()) + path = r"C:\Users\someone\git\PyRIT\dbdata\prompt-memory-entries\media\1780010098266691.bin" + piece = MessagePiece( + role="assistant", + original_value=path, + original_value_data_type=data_type, + converted_value=path, + converted_value_data_type=data_type, + conversation_id=conv_id, + sequence=0, + ) + sqlite_instance._insert_entry(PromptMemoryEntry(entry=piece)) + + result = sqlite_instance.get_conversation_stats(conversation_ids=[conv_id]) + stats = result[conv_id] + + assert stats.last_message_data_type == data_type + # Memory returns the raw value (truncated up to PREVIEW_FETCH_MAX_LEN); + # formatting/labeling is the backend mapper's responsibility. + assert stats.last_message_preview == path + + +def test_get_conversation_stats_uses_last_piece_data_type(sqlite_instance): + """Stats reflect the data type of the most recent message, not the + first one, so the backend mapper picks the right rendering.""" + import uuid + + from pyrit.models import MessagePiece + + conv_id = str(uuid.uuid4()) + text_piece = MessagePiece( + role="user", + original_value="hi there", + original_value_data_type="text", + converted_value="hi there", + converted_value_data_type="text", + conversation_id=conv_id, + sequence=0, + ) + audio_path = r"C:\dbdata\prompt-memory-entries\audio\response.mp3" + media_piece = MessagePiece( + role="assistant", + original_value=audio_path, + original_value_data_type="audio_path", + converted_value=audio_path, + converted_value_data_type="audio_path", + conversation_id=conv_id, + sequence=1, + ) + sqlite_instance._insert_entries(entries=[PromptMemoryEntry(entry=text_piece), PromptMemoryEntry(entry=media_piece)]) + + result = sqlite_instance.get_conversation_stats(conversation_ids=[conv_id]) + stats = result[conv_id] + + assert stats.last_message_data_type == "audio_path" + assert stats.last_message_preview == audio_path + + def test_dispose_engine_tolerates_closed_log_stream(sqlite_instance, capsys): """Verify dispose_engine does not raise or emit 'Logging error' when streams are closed (GH-1520).""" pyrit_logger = logging.getLogger("pyrit") diff --git a/tests/unit/models/test_conversation_stats.py b/tests/unit/models/test_conversation_stats.py index 44a58a82f9..f7bb64d85d 100644 --- a/tests/unit/models/test_conversation_stats.py +++ b/tests/unit/models/test_conversation_stats.py @@ -13,6 +13,7 @@ def test_conversation_stats_defaults(): stats = ConversationStats() assert stats.message_count == 0 assert stats.last_message_preview is None + assert stats.last_message_data_type is None assert stats.labels == {} assert stats.created_at is None @@ -22,11 +23,13 @@ def test_conversation_stats_with_values(): stats = ConversationStats( message_count=5, last_message_preview="Hello world", + last_message_data_type="text", labels={"env": "test"}, created_at=now, ) assert stats.message_count == 5 assert stats.last_message_preview == "Hello world" + assert stats.last_message_data_type == "text" assert stats.labels == {"env": "test"} assert stats.created_at == now @@ -41,6 +44,27 @@ def test_conversation_stats_preview_max_len_class_var(): assert ConversationStats.PREVIEW_MAX_LEN == 100 +def test_conversation_stats_preview_fetch_max_len_class_var(): + # Storage-fetch cap must be strictly larger than the display-truncation + # cap so downstream formatters can extract a basename / preview from a + # long media path before further truncation. + assert ConversationStats.PREVIEW_FETCH_MAX_LEN > ConversationStats.PREVIEW_MAX_LEN + + +def test_conversation_stats_accepts_media_data_type(): + stats = ConversationStats( + message_count=1, + last_message_preview=r"C:\foo\bar.png", + last_message_data_type="image_path", + ) + assert stats.last_message_data_type == "image_path" + + +def test_conversation_stats_rejects_unknown_data_type(): + with pytest.raises(ValidationError): + ConversationStats(last_message_data_type="not_a_real_type") # type: ignore[arg-type] + + def test_conversation_stats_labels_default_factory(): stats1 = ConversationStats() stats2 = ConversationStats()