diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 080fe6bde8..10775019d2 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -737,7 +737,10 @@ def _execute_get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> p2.converted_value_data_type AS last_data_type FROM "PromptMemoryEntries" p2 WHERE p2.conversation_id = aggregate_rows.conversation_id - ORDER BY p2.sequence DESC, p2.id DESC + ORDER BY + p2.sequence DESC, + p2.timestamp DESC, + LOWER(CAST(p2.id AS VARCHAR(36))) COLLATE Latin1_General_100_BIN2 DESC ) AS latest """ ).columns(created_at=UTCDateTime()) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 1756c579c4..ef5cf2e114 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -800,7 +800,7 @@ def _execute_get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> SELECT p2.id FROM "PromptMemoryEntries" p2 WHERE p2.conversation_id = aggregate_rows.conversation_id - ORDER BY p2.sequence DESC, p2.id DESC + ORDER BY p2.sequence DESC, p2.timestamp DESC, p2.id DESC LIMIT 1 ) """ diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index a2e5b8a8bf..dc445517fc 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -621,6 +621,27 @@ def test_get_conversation_stats_types_created_at_as_utc_datetime( assert isinstance(sql.selected_columns.created_at.type, UTCDateTime) +def test_get_conversation_stats_orders_the_last_piece_canonically( + uninitialized_memory_interface: AzureSQLMemory, +) -> None: + """The last message lookup orders pieces the way the conversation page does. + + SQL Server UUID ordering differs from the SDK string comparison the page + uses, so the lookup sorts ids through the same lowered collation, and the + timestamp decides between pieces that share a sequence number. + """ + session = MagicMock() + session.execute.return_value.fetchall.return_value = [] + + with patch.object(uninitialized_memory_interface, "_get_session", return_value=session): + uninitialized_memory_interface._execute_get_conversation_stats(conversation_ids=["conversation"]) + + sql = str(session.execute.call_args.args[0]) + assert "p2.sequence DESC" in sql + assert "p2.timestamp DESC" in sql + assert "LOWER(CAST(p2.id AS VARCHAR(36))) COLLATE Latin1_General_100_BIN2 DESC" in sql + + def test_get_conversation_stats_uses_one_latest_row_apply( uninitialized_memory_interface: AzureSQLMemory, ) -> None: diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index 1590718dde..1e7ab50719 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -845,7 +845,7 @@ def capture_statement(conn, cursor, statement, parameters, context, executemany) sql = "\n".join(statements).upper() assert 'LEFT JOIN "PROMPTMEMORYENTRIES" LATEST' in sql - assert "ORDER BY P2.SEQUENCE DESC, P2.ID DESC" in sql + assert "ORDER BY P2.SEQUENCE DESC, P2.TIMESTAMP DESC, P2.ID DESC" in sql assert "ROW_NUMBER" not in sql @@ -928,6 +928,52 @@ async def test_get_conversation_stats_does_not_read_labels_from_prompt_entries(s assert result[conv_id].labels == {} +async def test_get_conversation_stats_picks_the_canonical_last_piece(sqlite_instance): + """The preview names the piece the conversation page orders last. + + Pieces of one turn share their sequence number, and the page orders them + by timestamp before the id tiebreak, so the preview lookup has to honor + the same order instead of handing the win to the bigger id. + """ + import uuid + from datetime import UTC, datetime, timedelta + + from pyrit.models import MessagePiece + + conv_id = str(uuid.uuid4()) + base = datetime(2026, 1, 1, tzinfo=UTC) + earlier = MessagePiece( + id=uuid.UUID("ffffffff-0000-0000-0000-000000000001"), + role="assistant", + original_value="earlier piece", + original_value_data_type="text", + converted_value="earlier piece", + converted_value_data_type="text", + conversation_id=conv_id, + sequence=2, + timestamp=base, + ) + later = MessagePiece( + id=uuid.UUID("00000000-0000-0000-0000-000000000002"), + role="assistant", + original_value="later piece", + original_value_data_type="text", + converted_value="later piece", + converted_value_data_type="text", + conversation_id=conv_id, + sequence=2, + timestamp=base + timedelta(seconds=10), + ) + for piece in (earlier, later): + sqlite_instance._insert_entry(PromptMemoryEntry(entry=piece)) + + stats = await sqlite_instance.get_conversation_stats_async(conversation_ids=[conv_id]) + page = await sqlite_instance.get_conversation_message_page_async(conversation_id=conv_id, limit=100) + + assert stats[conv_id].last_message_preview == "later piece" + assert page.messages[-1].message_pieces[-1].converted_value == "later piece" + + async 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.