From 269b4f93f7f55bd8be49bff59f27131b0c5ec6aa Mon Sep 17 00:00:00 2001 From: behnamousat Date: Fri, 10 Jul 2026 15:41:26 -0700 Subject: [PATCH 01/24] Persist target identifiers as a deduped, content-addressed store --- ...f7a9c1b3d2_add_target_identifiers_table.py | 203 ++++++++++++++++++ pyrit/memory/memory_interface.py | 60 +++++- pyrit/memory/memory_models.py | 194 ++++++++++++++++- .../test_interface_prompts.py | 96 +++++++++ 4 files changed, 551 insertions(+), 2 deletions(-) create mode 100644 pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py new file mode 100644 index 0000000000..f64af89ef6 --- /dev/null +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py @@ -0,0 +1,203 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Introduce the TargetIdentifiers table and reference it from Conversations. + +Phase 1 (dual-write) of storing component identifiers as first-class, +content-addressed rows. Creates ``TargetIdentifiers`` (one row per distinct +target identifier, keyed by its content ``hash``, with promoted scalar query +columns) and ``TargetIdentifierChildren`` (a self-referential pivot mapping a +multi-target to its inner target identifiers), adds a nullable +``target_identifier_hash`` foreign key to ``Conversations``, and backfills all +three from the existing ``Conversations.target_identifier`` JSON column. The JSON +column is retained (reads still come from it), so this migration is purely +additive. + +Revision ID: e5f7a9c1b3d2 +Revises: d4e6f8a0b2c4 +Create Date: 2026-07-10 12:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "e5f7a9c1b3d2" +down_revision: str | None = "d4e6f8a0b2c4" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "TargetIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("endpoint", sa.String(), nullable=True), + sa.Column("model_name", sa.String(), nullable=True), + sa.Column("underlying_model_name", sa.String(), nullable=True), + sa.Column("temperature", sa.Float(), nullable=True), + sa.Column("top_p", sa.Float(), nullable=True), + sa.Column("max_requests_per_minute", sa.Integer(), nullable=True), + sa.Column("supported_auth_modes", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + # Self-referential pivot mapping a multi-target to its inner target identifiers. + # Both endpoints are content hashes into TargetIdentifiers; ``position`` preserves + # the parent's ``targets`` list order. Named FK constraints for SQL Server / batch + # portability. + op.create_table( + "TargetIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_child_hash" + ), + ) + + # Batch op for SQLite portability (no ALTER TABLE ADD FOREIGN KEY on SQLite). + # The FK constraint must be named explicitly: Alembic batch mode rejects an + # unnamed constraint. + with op.batch_alter_table("Conversations") as batch_op: + batch_op.add_column(sa.Column("target_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_conversations_target_identifier_hash", + "TargetIdentifiers", + ["target_identifier_hash"], + ["hash"], + ) + + _backfill_target_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("Conversations") as batch_op: + batch_op.drop_column("target_identifier_hash") + # Drop the child edge table before its referenced parent table. + op.drop_table("TargetIdentifierChildren") + op.drop_table("TargetIdentifiers") + + +def _backfill_target_identifiers() -> None: + """ + Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set + ``Conversations.target_identifier_hash``. + + For every ``Conversations`` row with a non-null ``target_identifier`` JSON, + reconstruct the ``TargetIdentifier`` (recomputing its content hash), insert the + deduped ``TargetIdentifiers`` row if absent -- recursing into any inner + ``targets`` first so the child edge foreign keys resolve -- record the + ``parent_hash -> child_hash`` edges, and point the conversation's + ``target_identifier_hash`` at the top-level row. Idempotent: hashes already present + are not re-inserted. Rows whose stored target cannot be reconstructed are logged and + skipped rather than aborting the upgrade. + """ + from pyrit.models import TargetIdentifier + + bind = op.get_bind() + rows = bind.execute( + sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') + ).fetchall() + + existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} + + insert_stmt = sa.text( + 'INSERT INTO "TargetIdentifiers" ' + "(hash, class_name, class_module, identifier_json, endpoint, model_name, " + "underlying_model_name, temperature, top_p, max_requests_per_minute, " + "supported_auth_modes, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " + ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " + ":supported_auth_modes, :pyrit_version)" + ) + edge_stmt = sa.text( + 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') + + def _insert_target(identifier: TargetIdentifier) -> int: + """ + Insert ``identifier`` and its descendants if absent; record child edges. + + Returns: + int: The number of new ``TargetIdentifiers`` rows inserted (self + descendants). + """ + target_hash = identifier.hash + if target_hash in existing_hashes: + return 0 + # Children first so the parent's edge foreign keys resolve. + inserted = 0 + for child in identifier.targets: + inserted += _insert_target(child) + bind.execute( + insert_stmt, + { + "hash": target_hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "endpoint": identifier.endpoint, + "model_name": identifier.model_name, + "underlying_model_name": identifier.underlying_model_name, + "temperature": identifier.temperature, + "top_p": identifier.top_p, + "max_requests_per_minute": identifier.max_requests_per_minute, + "supported_auth_modes": ( + json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None + ), + "pyrit_version": identifier.pyrit_version, + }, + ) + existing_hashes.add(target_hash) + inserted += 1 + for position, child in enumerate(identifier.targets): + bind.execute(edge_stmt, {"parent_hash": target_hash, "position": position, "child_hash": child.hash}) + return inserted + + inserted = 0 + linked = 0 + skipped = 0 + for conversation_id, raw_target in rows: + stored = json.loads(raw_target) if isinstance(raw_target, str) else raw_target + if not stored: + continue + try: + identifier = TargetIdentifier.model_validate(stored) + except Exception: + skipped += 1 + logger.warning( + f"TargetIdentifiers backfill: could not reconstruct target for " + f"conversation_id {conversation_id!r}; skipping." + ) + continue + + inserted += _insert_target(identifier) + bind.execute(update_stmt, {"hash": identifier.hash, "cid": conversation_id}) + linked += 1 + + if inserted or linked or skipped: + logger.info( + f"TargetIdentifiers backfill: inserted {inserted} identifier row(s); " + f"linked {linked} conversation(s); skipped {skipped}." + ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 9e2bf380af..7ff7943c21 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -14,7 +14,7 @@ from sqlalchemy import MetaData, and_, not_, or_, select from sqlalchemy.engine.base import Engine -from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.orm.attributes import InstrumentedAttribute if TYPE_CHECKING: @@ -29,6 +29,8 @@ ScenarioResultEntry, ScoreEntry, SeedEntry, + TargetIdentifierChildEntry, + TargetIdentifierEntry, ) from pyrit.memory.storage import ( DataTypeSerializer, @@ -50,6 +52,7 @@ SeedDataset, SeedGroup, SeedType, + TargetIdentifier, group_conversation_message_pieces_by_sequence, sort_message_pieces, ) @@ -459,6 +462,13 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: try: existing = session.get(ConversationEntry, conversation.conversation_id) if existing is None: + if conversation.target_identifier is not None: + self._persist_target_identifier( + session=session, + target_identifier=TargetIdentifier.from_component_identifier( + conversation.target_identifier + ), + ) session.add(entry) elif ( entry.target_identifier is not None @@ -476,6 +486,54 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}") raise + @staticmethod + def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentifier) -> None: + """ + Persist ``target_identifier`` and its inner targets as content-addressed rows. + + Mirrors the conversation-registration contract ("ensure dependencies exist + first"): every inner target in ``target_identifier.targets`` is persisted + recursively before this identifier's own row, so the ``TargetIdentifierChildren`` + edges (which foreign-key both endpoints into ``TargetIdentifiers``) resolve. + Identifier rows are immutable and keyed by their content hash, so this is an + idempotent insert-if-absent: an identical target -- or a shared inner target -- + reused across many conversations maps to a single row. Runs inside the caller's + session so every row is flushed before the FK-bearing ``ConversationEntry`` is + added. + + If the row already exists it was fully persisted before (children and edges + included, since rows are immutable), so this returns early. Otherwise the row and + its child edges are inserted inside a savepoint: a concurrent writer that inserts + the same hash first surfaces as an ``IntegrityError`` that rolls back only this + insert (not the surrounding write) and is then treated as a no-op -- the row + exists either way. + + Args: + session (Any): The active SQLAlchemy session (the caller's transaction). + target_identifier (TargetIdentifier): The target identifier to persist. + """ + if session.get(TargetIdentifierEntry, target_identifier.hash) is not None: + return + # Children first, so the parent's edge foreign keys resolve on flush. + for inner_target in target_identifier.targets: + MemoryInterface._persist_target_identifier(session=session, target_identifier=inner_target) + try: + with session.begin_nested(): + session.add(TargetIdentifierEntry.from_domain_model(target_identifier)) + for position, inner_target in enumerate(target_identifier.targets): + session.add( + TargetIdentifierChildEntry( + parent_hash=target_identifier.hash, + position=position, + child_hash=inner_target.hash, + ) + ) + session.flush() + except IntegrityError: + # A racing writer inserted the same content-addressed row; immutable + # rows make this a safe no-op. + pass + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a7cfdbf552..04a0fb568e 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -4,9 +4,10 @@ import json import logging import uuid +from abc import abstractmethod from collections.abc import Sequence from datetime import datetime, timezone -from typing import Any, Literal +from typing import Any, Generic, Literal, TypeVar from pydantic import BaseModel, ConfigDict from sqlalchemy import ( @@ -29,6 +30,7 @@ relationship, ) from sqlalchemy.types import Uuid +from typing_extensions import Self import pyrit from pyrit.common.utils import to_sha256 @@ -56,6 +58,7 @@ SeedPrompt, SeedSimulatedConversation, SeedType, + TargetIdentifier, ) logger = logging.getLogger(__name__) @@ -353,6 +356,185 @@ def __str__(self) -> str: return f"{self.role}: {self.converted_value}" +TDomain = TypeVar("TDomain") + + +class DomainBackedEntry(Base, Generic[TDomain]): + """ + Mixin marking a DB entry as the persistence representation of a domain model. + + Every ``*Entry`` in this module mirrors a domain model (``PromptMemoryEntry`` and + ``MessagePiece``, ``ScoreEntry`` and ``Score``, ``TargetIdentifierEntry`` and + ``TargetIdentifier``, and so on). ``from_domain_model`` is the single, uniform seam + that converts a domain model into an unsaved row, so the domain-to-DB direction has + one well-known name across every entry that adopts this base. + """ + + __abstract__ = True + + @classmethod + @abstractmethod + def from_domain_model(cls, domain_model: TDomain) -> Self: + """ + Build an unsaved entry row from its domain model. + + Args: + domain_model (TDomain): The domain model this entry persists. + + Returns: + Self: A new, unsaved row. + """ + + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Reject concrete subclasses that do not implement ``from_domain_model``. + + The SQLAlchemy declarative ``Base`` is not an ``ABCMeta``, so a bare + ``@abstractmethod`` would not stop a concrete entry from omitting the converter. + This fires at class-definition time so a dev who forgets is told immediately. + SQLAlchemy abstract/intermediate mapped classes (``__abstract__ = True``) are + skipped so they can leave the method abstract for their concrete subclasses. + + Raises: + TypeError: If a concrete (non-``__abstract__``) subclass leaves + ``from_domain_model`` abstract. + """ + super().__init_subclass__(**kwargs) + if cls.__dict__.get("__abstract__", False): + return + method = getattr(cls, "from_domain_model", None) + if method is None or getattr(method, "__isabstractmethod__", False): + raise TypeError( + f"{cls.__name__} inherits DomainBackedEntry but does not implement " + "from_domain_model(...); every concrete entry must define how its " + "domain model is converted into a row." + ) + + +T = TypeVar("T", bound=ComponentIdentifier) + + +class ComponentIdentifierEntry(DomainBackedEntry[T]): + """ + Abstract base for tables that persist a ``ComponentIdentifier`` projection. + + Mirrors the identifier class hierarchy: concrete identifier tables inherit the + shared, always-populated columns the way ``TargetIdentifier`` inherits + ``ComponentIdentifier``. The content ``hash`` is the natural, dedupable primary + key, and ``identifier_json`` holds the full flat dump for lossless + reconstruction (children stay inline). Rows are immutable — the same content + always maps to the same hash, so a given identifier reused across rows is + stored once. + + Subclasses declare their promoted query columns and implement the + ``DomainBackedEntry.from_domain_model`` seam to map their strongly-typed identifier + projection onto the shared columns plus those promoted columns. The shared columns + are built once, here, so subclasses never repeat that logic. + """ + + __abstract__ = True + + #: Content-addressed identity — the same value as ``ComponentIdentifier.hash``. + #: SHA256 hex digest is 64 chars; bounded for SQL Server key/index compatibility. + hash: Mapped[str] = mapped_column(String(64), primary_key=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + class_module: Mapped[str] = mapped_column(String, nullable=False) + #: Full flat ``model_dump()`` of the identifier. Source of truth on reload. + identifier_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) + #: Version that first wrote this content-addressed row. Nullable for backwards + #: compatibility with existing databases. + pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) + + +class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): + """ + Content-addressed store of ``TargetIdentifier`` projections, deduped by hash. + + Populated as a side effect of registering a conversation (see + ``MemoryInterface._persist_target_identifier``). ``ConversationEntry`` references a + row here via ``target_identifier_hash``. The promoted scalar columns (``endpoint`` / + ``model_name`` / ``underlying_model_name`` / ``temperature`` / ``top_p`` / + ``max_requests_per_minute`` / ``supported_auth_modes``) are surfaced for querying; + ``identifier_json`` remains the source of truth on reload. Inner targets of a + multi-target are linked via ``TargetIdentifierChildren`` (and also live inline in + ``identifier_json``). + """ + + __tablename__ = "TargetIdentifiers" + __table_args__ = {"extend_existing": True} + + endpoint: Mapped[str | None] = mapped_column(String, nullable=True) + model_name: Mapped[str | None] = mapped_column(String, nullable=True) + underlying_model_name: Mapped[str | None] = mapped_column(String, nullable=True) + temperature: Mapped[float | None] = mapped_column(Float, nullable=True) + top_p: Mapped[float | None] = mapped_column(Float, nullable=True) + max_requests_per_minute: Mapped[int | None] = mapped_column(INTEGER, nullable=True) + supported_auth_modes: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + + @classmethod + def from_domain_model(cls, domain_model: TargetIdentifier) -> Self: + """ + Build a row from a ``TargetIdentifier``. + + Maps the shared columns plus the promoted scalar query columns; the full flat + dump (including any inner ``targets``) is stored in ``identifier_json``. Inner + target edges are written separately by + ``MemoryInterface._persist_target_identifier``. + + Args: + domain_model (TargetIdentifier): The target identifier to persist. + + Returns: + Self: A new, unsaved row keyed by the identifier's content hash. + """ + return cls( + hash=domain_model.hash, + class_name=domain_model.class_name, + class_module=domain_model.class_module, + identifier_json=domain_model.model_dump(), + pyrit_version=domain_model.pyrit_version, + endpoint=domain_model.endpoint, + model_name=domain_model.model_name, + underlying_model_name=domain_model.underlying_model_name, + temperature=domain_model.temperature, + top_p=domain_model.top_p, + max_requests_per_minute=domain_model.max_requests_per_minute, + supported_auth_modes=domain_model.supported_auth_modes, + ) + + +class TargetIdentifierChildEntry(Base): + """ + Ordered edge rows linking a multi-target ``TargetIdentifierEntry`` to its inner + target identifiers. + + A multi-target (e.g. ``RoundRobinTarget``) owns a ``targets`` list; each inner + target is itself a content-addressed ``TargetIdentifiers`` row, and one edge row + here maps ``parent_hash -> child_hash`` at a given ``position`` (the child's index + in the parent's ``targets`` list). Both endpoints are hashes into + ``TargetIdentifiers``, so an inner target shared across parents dedupes to a single + row and is merely referenced here. This is a query index over target composition; + ``TargetIdentifierEntry.identifier_json`` remains the source of truth for + reconstruction (inner targets are stored inline there too). + + Materialized by ``MemoryInterface._persist_target_identifier`` while persisting a + target and its children; it is never built from a single domain model, so it is a + plain ``Base`` row rather than a ``DomainBackedEntry``. + """ + + __tablename__ = "TargetIdentifierChildren" + __table_args__ = {"extend_existing": True} + + parent_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + #: Zero-based index of the child within the parent's ``targets`` list. + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + child_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=False + ) + + class ConversationEntry(Base): """ Conversation-scoped metadata, persisted once per ``conversation_id``. @@ -362,6 +544,10 @@ class ConversationEntry(Base): row. The target is captured once when the conversation's pieces are written and read back via ``MemoryInterface._get_conversation`` (it is not stamped onto individual pieces). + + The target is dual-written: the full identifier stays in the ``target_identifier`` + JSON column (still the read source), and ``target_identifier_hash`` references the + deduped ``TargetIdentifierEntry`` row keyed by the identifier's content hash. """ __tablename__ = "Conversations" @@ -369,6 +555,11 @@ class ConversationEntry(Base): conversation_id = mapped_column(String(36), primary_key=True, nullable=False) target_identifier: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True) + #: Foreign key to the content-addressed ``TargetIdentifiers`` row. Nullable: + #: a conversation may be registered without a known target. + target_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) # Version of PyRIT used when this entry was created. Nullable for backwards # compatibility with existing databases. @@ -383,6 +574,7 @@ def __init__(self, *, conversation: Conversation) -> None: """ self.conversation_id = conversation.conversation_id self.target_identifier = conversation.target_identifier.model_dump() if conversation.target_identifier else None + self.target_identifier_hash = conversation.target_identifier.hash if conversation.target_identifier else None self.pyrit_version = pyrit.__version__ def get_conversation(self) -> Conversation: diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index b3860519e3..3a50b09ceb 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -679,6 +679,102 @@ def test_add_conversation_to_memory_different_target_reregister_raises(sqlite_in assert metadata.target_identifier.hash == target_a.hash +def test_target_identifier_dual_write_reconstruction_is_equivalent(sqlite_instance: MemoryInterface): + # Phase 1 dual-write invariant: reconstructing the target from the ConversationEntry + # JSON column must be identical to reconstructing it from the deduped + # TargetIdentifierEntry.identifier_json row, and the stored PK must match the + # recomputed content hash. This is the safety net that lets phase 2 flip reads to + # the FK path. + from pyrit.memory.memory_models import ConversationEntry, TargetIdentifierEntry + + target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://api.openai.com", "model_name": "gpt-4"}, + ) + conversation_id = "conv-dualwrite" + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id=conversation_id, target_identifier=target) + ) + + conv_entry = sqlite_instance._query_entries( + ConversationEntry, conditions=ConversationEntry.conversation_id == conversation_id + )[0] + id_entry = sqlite_instance._query_entries( + TargetIdentifierEntry, conditions=TargetIdentifierEntry.hash == target.hash + )[0] + + from_conversation_json = ComponentIdentifier.model_validate(conv_entry.target_identifier) + from_identifier_row = ComponentIdentifier.model_validate(id_entry.identifier_json) + + # Both reconstructions agree (equality is content-hash based) and match the FK / PK. + assert from_conversation_json == from_identifier_row + assert from_conversation_json.hash == target.hash + assert id_entry.hash == target.hash + assert conv_entry.target_identifier_hash == target.hash + # Promoted columns are surfaced from params for querying. + assert id_entry.endpoint == "https://api.openai.com" + assert id_entry.model_name == "gpt-4" + + +def test_target_identifier_row_is_deduped_across_conversations(sqlite_instance: MemoryInterface): + # The same target reused across conversations is content-addressed, so it is stored + # once: two conversations with an identical target share a single TargetIdentifiers row. + from pyrit.memory.memory_models import TargetIdentifierEntry + + target = ComponentIdentifier( + class_name="OpenAIChatTarget", class_module="pyrit.prompt_target", params={"endpoint": "shared"} + ) + for cid in ("conv-dedup-a", "conv-dedup-b"): + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id=cid, target_identifier=target) + ) + + rows = sqlite_instance._query_entries(TargetIdentifierEntry, conditions=TargetIdentifierEntry.hash == target.hash) + assert len(rows) == 1 + + +def test_target_identifier_persists_inner_targets_and_edges(sqlite_instance: MemoryInterface): + # A multi-target's inner targets are persisted as their own content-addressed rows + # and linked to the parent via ordered TargetIdentifierChildren edges. Promoted + # scalar columns are surfaced on each row for querying. + from pyrit.memory.memory_models import TargetIdentifierChildEntry, TargetIdentifierEntry + + inner_a = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://a", "model_name": "gpt-a", "temperature": 0.5}, + ) + inner_b = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://b", "model_name": "gpt-b"}, + ) + multi = ComponentIdentifier( + class_name="RoundRobinTarget", + class_module="pyrit.prompt_target", + children={"targets": [inner_a, inner_b]}, + ) + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id="conv-multi", target_identifier=multi) + ) + + id_rows = sqlite_instance._query_entries(TargetIdentifierEntry) + by_hash = {row.hash: row for row in id_rows} + # Parent + both inner targets are each stored once (content-addressed). + assert multi.hash in by_hash + assert inner_a.hash in by_hash + assert inner_b.hash in by_hash + # Promoted scalar column surfaced from the inner target's params. + assert by_hash[inner_a.hash].temperature == 0.5 + + edges = sqlite_instance._query_entries( + TargetIdentifierChildEntry, conditions=TargetIdentifierChildEntry.parent_hash == multi.hash + ) + ordered = sorted(edges, key=lambda edge: edge.position) + assert [(edge.position, edge.child_hash) for edge in ordered] == [(0, inner_a.hash), (1, inner_b.hash)] + + def test_insert_conversation_rolls_back_and_reraises_on_db_error(sqlite_instance: MemoryInterface): # A DB failure during registration rolls back the session and propagates the error # rather than leaving a half-written Conversations row. From 530837b9fcf8f8600796bf0068e0364c4218b5d8 Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 12:31:41 -0700 Subject: [PATCH 02/24] reuse promoted fields --- pyrit/memory/memory_models.py | 142 ++++++++++++++---- .../identifiers/component_identifier.py | 29 ++++ 2 files changed, 139 insertions(+), 32 deletions(-) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 04a0fb568e..cfb9f11af4 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -5,9 +5,10 @@ import logging import uuid from abc import abstractmethod -from collections.abc import Sequence +from collections.abc import Callable, Sequence +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Generic, Literal, TypeVar +from typing import Any, ClassVar, Generic, Literal, TypeVar, cast from pydantic import BaseModel, ConfigDict from sqlalchemy import ( @@ -414,6 +415,16 @@ def __init_subclass__(cls, **kwargs: Any) -> None: T = TypeVar("T", bound=ComponentIdentifier) +@dataclass(frozen=True) +class _ChildRelationshipSpec: + """Mapping from a promoted child field to its ORM edge relationship wiring.""" + + relationship_name: str + edge_factory: Callable[[], Any] + edge_child_attr: str + edge_position_attr: str | None = "position" + + class ComponentIdentifierEntry(DomainBackedEntry[T]): """ Abstract base for tables that persist a ``ComponentIdentifier`` projection. @@ -434,6 +445,10 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): __abstract__ = True + #: Optional per-child-field wiring for materialized edge relationships. + #: Empty by default so identifier rows only persist their own projection. + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = {} + #: Content-addressed identity — the same value as ``ComponentIdentifier.hash``. #: SHA256 hex digest is 64 chars; bounded for SQL Server key/index compatibility. hash: Mapped[str] = mapped_column(String(64), primary_key=True) @@ -445,6 +460,68 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): #: compatibility with existing databases. pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) + @classmethod + def from_domain_model(cls, domain_model: T) -> Self: + """ + Build an unsaved component identifier memory entry from the given domain model. + + Args: + domain_model (T): The domain model this entry persists. + + Returns: + Self: A new, unsaved row. + """ + return cls._from_domain_model_recursive(domain_model=domain_model, seen={}) + + @classmethod + def _from_domain_model_recursive(cls, *, domain_model: T, seen: dict[str, Self]) -> Self: + existing = seen.get(domain_model.hash) + if existing is not None: + return existing + + entry = cls._from_domain_model_shallow(domain_model=domain_model) + seen[domain_model.hash] = entry + cls._attach_child_relationship_rows(entry=entry, domain_model=domain_model, seen=seen) + return entry + + @classmethod + def _from_domain_model_shallow(cls, *, domain_model: T) -> Self: + entry = cls( + hash=domain_model.hash, + class_name=domain_model.class_name, + class_module=domain_model.class_module, + identifier_json=domain_model.model_dump(), + pyrit_version=domain_model.pyrit_version, + ) + for name, value in domain_model.promoted_scalar_values().items(): + setattr(entry, name, value) # each promoted scalar → its mapped column + return entry + + @classmethod + def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: dict[str, Self]) -> None: + for field_name in domain_model.promoted_child_field_names(): + spec = cls.CHILD_RELATIONSHIP_SPECS.get(field_name) + if spec is None: + continue + + child_value = getattr(domain_model, field_name) + children = ( + child_value + if isinstance(child_value, list) + else ([child_value] if child_value is not None else []) + ) + edge_rows = getattr(entry, spec.relationship_name) + for position, child_identifier in enumerate(children): + if not isinstance(child_identifier, ComponentIdentifier): + continue + typed_child_identifier = cast("T", child_identifier) + child_entry = cls._from_domain_model_recursive(domain_model=typed_child_identifier, seen=seen) + edge_row = spec.edge_factory() + setattr(edge_row, spec.edge_child_attr, child_entry) + if spec.edge_position_attr: + setattr(edge_row, spec.edge_position_attr, position) + edge_rows.append(edge_row) + class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): """ @@ -463,6 +540,15 @@ class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): __tablename__ = "TargetIdentifiers" __table_args__ = {"extend_existing": True} + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = { + "targets": _ChildRelationshipSpec( + relationship_name="targets", + edge_factory=lambda: TargetIdentifierChildEntry(), + edge_child_attr="child", + edge_position_attr="position", + ) + } + endpoint: Mapped[str | None] = mapped_column(String, nullable=True) model_name: Mapped[str | None] = mapped_column(String, nullable=True) underlying_model_name: Mapped[str | None] = mapped_column(String, nullable=True) @@ -471,36 +557,16 @@ class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): max_requests_per_minute: Mapped[int | None] = mapped_column(INTEGER, nullable=True) supported_auth_modes: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) - @classmethod - def from_domain_model(cls, domain_model: TargetIdentifier) -> Self: - """ - Build a row from a ``TargetIdentifier``. - - Maps the shared columns plus the promoted scalar query columns; the full flat - dump (including any inner ``targets``) is stored in ``identifier_json``. Inner - target edges are written separately by - ``MemoryInterface._persist_target_identifier``. - - Args: - domain_model (TargetIdentifier): The target identifier to persist. - - Returns: - Self: A new, unsaved row keyed by the identifier's content hash. - """ - return cls( - hash=domain_model.hash, - class_name=domain_model.class_name, - class_module=domain_model.class_module, - identifier_json=domain_model.model_dump(), - pyrit_version=domain_model.pyrit_version, - endpoint=domain_model.endpoint, - model_name=domain_model.model_name, - underlying_model_name=domain_model.underlying_model_name, - temperature=domain_model.temperature, - top_p=domain_model.top_p, - max_requests_per_minute=domain_model.max_requests_per_minute, - supported_auth_modes=domain_model.supported_auth_modes, - ) + #: Ordered child-edge rows (``parent_hash -> child_hash``) for this target. + #: Reconstructing nested targets from relational rows requires joining through + #: this relationship and then following ``TargetIdentifierChildEntry.child``. + targets: Mapped[list["TargetIdentifierChildEntry"]] = relationship( + "TargetIdentifierChildEntry", + primaryjoin="TargetIdentifierEntry.hash == TargetIdentifierChildEntry.parent_hash", + foreign_keys="TargetIdentifierChildEntry.parent_hash", + order_by="TargetIdentifierChildEntry.position", + back_populates="parent", + ) class TargetIdentifierChildEntry(Base): @@ -534,6 +600,18 @@ class TargetIdentifierChildEntry(Base): String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=False ) + #: Parent target that owns this edge position. + parent: Mapped["TargetIdentifierEntry"] = relationship( + "TargetIdentifierEntry", + foreign_keys=[parent_hash], + back_populates="targets", + ) + #: Child target row referenced by ``child_hash``. + child: Mapped["TargetIdentifierEntry"] = relationship( + "TargetIdentifierEntry", + foreign_keys=[child_hash], + ) + class ConversationEntry(Base): """ diff --git a/pyrit/models/identifiers/component_identifier.py b/pyrit/models/identifiers/component_identifier.py index b577d4be25..8a49e5ce1b 100644 --- a/pyrit/models/identifiers/component_identifier.py +++ b/pyrit/models/identifiers/component_identifier.py @@ -345,6 +345,35 @@ def _promoted_child_fields(cls) -> tuple[str, ...]: """ return tuple(n for n in cls._promoted_fields() if cls._is_child_field(cls.model_fields[n].annotation)) + @classmethod + def promoted_scalar_field_names(cls) -> tuple[str, ...]: + """ + Get names of this identifier's promoted scalar (param) fields — the DB-column projection. + + Returns: + tuple[str, ...]: Promoted scalar field names. + """ + return cls._promoted_param_fields() + + @classmethod + def promoted_child_field_names(cls) -> tuple[str, ...]: + """ + Get names of this identifier's promoted child fields. + + Returns: + tuple[str, ...]: Promoted child field names. + """ + return cls._promoted_child_fields() + + def promoted_scalar_values(self) -> dict[str, Any]: + """ + Get this identifier's promoted scalar fields as ``{name: value}`` (children/targets excluded). + + Returns: + dict[str, Any]: Promoted scalar field names and values. + """ + return {name: getattr(self, name) for name in self._promoted_param_fields()} + @classmethod def get_reference_component_types(cls) -> dict[str, ComponentType]: """ From b0483f50735112a8617e868918bacb80e54c0972 Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 14:21:17 -0700 Subject: [PATCH 03/24] persist scorer identifiers --- ...c8e0f2b4d6_add_scorer_identifiers_table.py | 192 ++++++++++++++++++ pyrit/memory/memory_interface.py | 74 ++++--- pyrit/memory/memory_models.py | 81 +++++++- .../test_interface_prompts.py | 3 + .../memory_interface/test_interface_scores.py | 67 ++++++ tests/unit/memory/test_memory_models.py | 31 +++ tests/unit/memory/test_migration.py | 71 +++++++ 7 files changed, 490 insertions(+), 29 deletions(-) create mode 100644 pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py new file mode 100644 index 0000000000..184185e2f2 --- /dev/null +++ b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py @@ -0,0 +1,192 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist scorer identifiers as content-addressed rows. + +Revision ID: a6c8e0f2b4d6 +Revises: e5f7a9c1b3d2 +Create Date: 2026-07-13 12:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +revision: str = "a6c8e0f2b4d6" +down_revision: str | None = "e5f7a9c1b3d2" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "ScorerIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("scorer_type", sa.String(), nullable=True), + sa.Column("score_aggregator", sa.String(), nullable=True), + sa.Column("prompt_target_hash", sa.String(64), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["prompt_target_hash"], ["TargetIdentifiers.hash"], name="fk_scorer_identifiers_prompt_target_hash" + ), + ) + op.create_table( + "ScorerIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_child_hash" + ), + ) + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.add_column(sa.Column("scorer_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_score_entries_scorer_identifier_hash", + "ScorerIdentifiers", + ["scorer_identifier_hash"], + ["hash"], + ) + + _backfill_scorer_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.drop_column("scorer_identifier_hash") + op.drop_table("ScorerIdentifierChildren") + op.drop_table("ScorerIdentifiers") + + +def _backfill_scorer_identifiers() -> None: + """Backfill scorer rows and score foreign keys from the retained JSON column.""" + from pyrit.models import ScorerIdentifier, TargetIdentifier + + bind = op.get_bind() + score_rows = bind.execute( + sa.text( + 'SELECT id, scorer_class_identifier FROM "ScoreEntries" ' + "WHERE scorer_class_identifier IS NOT NULL" + ) + ).fetchall() + scorer_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScorerIdentifiers"')).fetchall()} + target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} + + target_insert = sa.text( + 'INSERT INTO "TargetIdentifiers" ' + "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " + "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " + ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " + ":supported_auth_modes, :pyrit_version)" + ) + target_edge_insert = sa.text( + 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + scorer_insert = sa.text( + 'INSERT INTO "ScorerIdentifiers" ' + "(hash, class_name, class_module, identifier_json, scorer_type, score_aggregator, " + "prompt_target_hash, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :scorer_type, " + ":score_aggregator, :prompt_target_hash, :pyrit_version)" + ) + scorer_edge_insert = sa.text( + 'INSERT INTO "ScorerIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') + + def _insert_target(identifier: TargetIdentifier) -> None: + if identifier.hash in target_hashes: + return + for child in identifier.targets: + _insert_target(child) + bind.execute( + target_insert, + { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "endpoint": identifier.endpoint, + "model_name": identifier.model_name, + "underlying_model_name": identifier.underlying_model_name, + "temperature": identifier.temperature, + "top_p": identifier.top_p, + "max_requests_per_minute": identifier.max_requests_per_minute, + "supported_auth_modes": ( + json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None + ), + "pyrit_version": identifier.pyrit_version, + }, + ) + target_hashes.add(identifier.hash) + for position, child in enumerate(identifier.targets): + bind.execute( + target_edge_insert, + {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, + ) + + def _insert_scorer(identifier: ScorerIdentifier) -> None: + if identifier.hash in scorer_hashes: + return + if identifier.prompt_target is not None: + _insert_target(identifier.prompt_target) + for child in identifier.sub_scorers: + _insert_scorer(child) + bind.execute( + scorer_insert, + { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "scorer_type": identifier.scorer_type, + "score_aggregator": identifier.score_aggregator, + "prompt_target_hash": identifier.prompt_target.hash if identifier.prompt_target is not None else None, + "pyrit_version": identifier.pyrit_version, + }, + ) + scorer_hashes.add(identifier.hash) + for position, child in enumerate(identifier.sub_scorers): + bind.execute( + scorer_edge_insert, + {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, + ) + + skipped = 0 + for score_id, raw_scorer in score_rows: + stored = json.loads(raw_scorer) if isinstance(raw_scorer, str) else raw_scorer + if not stored: + continue + try: + identifier = ScorerIdentifier.model_validate(stored) + _insert_scorer(identifier) + bind.execute(score_update, {"hash": identifier.hash, "id": score_id}) + except Exception: + skipped += 1 + logger.warning( + f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", + exc_info=True, + ) + + if skipped: + logger.warning(f"ScorerIdentifiers backfill skipped {skipped} score row(s)") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 7ff7943c21..b28dbd6ce8 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -28,8 +28,8 @@ PromptMemoryEntry, ScenarioResultEntry, ScoreEntry, + ScorerIdentifierEntry, SeedEntry, - TargetIdentifierChildEntry, TargetIdentifierEntry, ) from pyrit.memory.storage import ( @@ -48,6 +48,7 @@ MessagePiece, ScenarioResult, Score, + ScorerIdentifier, Seed, SeedDataset, SeedGroup, @@ -491,15 +492,11 @@ def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentif """ Persist ``target_identifier`` and its inner targets as content-addressed rows. - Mirrors the conversation-registration contract ("ensure dependencies exist - first"): every inner target in ``target_identifier.targets`` is persisted - recursively before this identifier's own row, so the ``TargetIdentifierChildren`` - edges (which foreign-key both endpoints into ``TargetIdentifiers``) resolve. - Identifier rows are immutable and keyed by their content hash, so this is an - idempotent insert-if-absent: an identical target -- or a shared inner target -- - reused across many conversations maps to a single row. Runs inside the caller's - session so every row is flushed before the FK-bearing ``ConversationEntry`` is - added. + The complete ORM graph is constructed from the domain model and merged into the + caller's session. SQLAlchemy reconciles shared child targets by primary key and + cascades the merge through ordered child edges. Identifier rows are immutable + and keyed by their content hash, so an identical target reused across many + conversations maps to a single row. If the row already exists it was fully persisted before (children and edges included, since rows are immutable), so this returns early. Otherwise the row and @@ -514,20 +511,10 @@ def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentif """ if session.get(TargetIdentifierEntry, target_identifier.hash) is not None: return - # Children first, so the parent's edge foreign keys resolve on flush. - for inner_target in target_identifier.targets: - MemoryInterface._persist_target_identifier(session=session, target_identifier=inner_target) try: with session.begin_nested(): - session.add(TargetIdentifierEntry.from_domain_model(target_identifier)) - for position, inner_target in enumerate(target_identifier.targets): - session.add( - TargetIdentifierChildEntry( - parent_hash=target_identifier.hash, - position=position, - child_hash=inner_target.hash, - ) - ) + entry = TargetIdentifierEntry.from_domain_model(target_identifier) + session.merge(entry) session.flush() except IntegrityError: # A racing writer inserted the same content-addressed row; immutable @@ -862,6 +849,9 @@ def add_scores_to_memory(self, *, scores: Sequence[Score]) -> None: Persisting the score even without a piece is intentional: aggregate analytics (e.g. refusal rate over a batch) still want the score row even when the scored content was never a real conversation turn. + + Raises: + SQLAlchemyError: If the score or identifier rows cannot be persisted. """ for score in scores: if score.message_piece_id: @@ -873,7 +863,45 @@ def add_scores_to_memory(self, *, scores: Sequence[Score]) -> None: # auto-link score to the original prompt id if the prompt is a duplicate if pieces[0].original_prompt_id != pieces[0].id: score.message_piece_id = pieces[0].original_prompt_id # type: ignore[ty:invalid-assignment] - self._insert_entries(entries=[ScoreEntry(entry=score) for score in scores]) + entries = [ScoreEntry(entry=score) for score in scores] + with closing(self.get_session()) as session: + try: + for entry in entries: + if entry.scorer_class_identifier: + self._persist_scorer_identifier( + session=session, + scorer_identifier=ScorerIdentifier.model_validate(entry.scorer_class_identifier), + ) + session.add_all(entries) + session.commit() + except SQLAlchemyError as e: + session.rollback() + logger.exception(f"Error inserting scores: {e}") + raise + + @staticmethod + def _persist_scorer_identifier(*, session: Any, scorer_identifier: ScorerIdentifier) -> None: + """Persist a complete scorer graph and its target dependencies.""" + if session.get(ScorerIdentifierEntry, scorer_identifier.hash) is not None: + return + + pending_scorers = [scorer_identifier] + while pending_scorers: + pending_scorer = pending_scorers.pop() + if pending_scorer.prompt_target is not None: + MemoryInterface._persist_target_identifier( + session=session, + target_identifier=pending_scorer.prompt_target, + ) + pending_scorers.extend(pending_scorer.sub_scorers) + + try: + with session.begin_nested(): + entry = ScorerIdentifierEntry.from_domain_model(scorer_identifier) + session.merge(entry) + session.flush() + except IntegrityError: + pass def get_scores( self, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index cfb9f11af4..ce81d3d75e 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -54,6 +54,7 @@ ScenarioRunState, Score, ScorerEvaluationIdentifier, + ScorerIdentifier, Seed, SeedObjective, SeedPrompt, @@ -506,9 +507,7 @@ def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: child_value = getattr(domain_model, field_name) children = ( - child_value - if isinstance(child_value, list) - else ([child_value] if child_value is not None else []) + child_value if isinstance(child_value, list) else ([child_value] if child_value is not None else []) ) edge_rows = getattr(entry, spec.relationship_name) for position, child_identifier in enumerate(children): @@ -583,9 +582,10 @@ class TargetIdentifierChildEntry(Base): ``TargetIdentifierEntry.identifier_json`` remains the source of truth for reconstruction (inner targets are stored inline there too). - Materialized by ``MemoryInterface._persist_target_identifier`` while persisting a - target and its children; it is never built from a single domain model, so it is a - plain ``Base`` row rather than a ``DomainBackedEntry``. + Constructed as part of the complete graph returned by + ``TargetIdentifierEntry.from_domain_model`` and persisted by + ``MemoryInterface._persist_target_identifier``. It is a plain ``Base`` row rather + than a ``DomainBackedEntry`` because an edge has no standalone domain model. """ __tablename__ = "TargetIdentifierChildren" @@ -613,6 +613,71 @@ class TargetIdentifierChildEntry(Base): ) +class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): + """Content-addressed store of ``ScorerIdentifier`` projections.""" + + __tablename__ = "ScorerIdentifiers" + __table_args__ = {"extend_existing": True} + + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = { + "sub_scorers": _ChildRelationshipSpec( + relationship_name="sub_scorers", + edge_factory=lambda: ScorerIdentifierChildEntry(), + edge_child_attr="child", + edge_position_attr="position", + ) + } + + scorer_type: Mapped[str | None] = mapped_column(String, nullable=True) + score_aggregator: Mapped[str | None] = mapped_column(String, nullable=True) + prompt_target_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) + + prompt_target: Mapped["TargetIdentifierEntry | None"] = relationship( + "TargetIdentifierEntry", + foreign_keys=[prompt_target_hash], + ) + sub_scorers: Mapped[list["ScorerIdentifierChildEntry"]] = relationship( + "ScorerIdentifierChildEntry", + primaryjoin="ScorerIdentifierEntry.hash == ScorerIdentifierChildEntry.parent_hash", + foreign_keys="ScorerIdentifierChildEntry.parent_hash", + order_by="ScorerIdentifierChildEntry.position", + back_populates="parent", + ) + + @classmethod + def _from_domain_model_shallow(cls, *, domain_model: ScorerIdentifier) -> Self: + entry = super()._from_domain_model_shallow(domain_model=domain_model) + entry.prompt_target_hash = domain_model.prompt_target.hash if domain_model.prompt_target is not None else None + return entry + + +class ScorerIdentifierChildEntry(Base): + """Ordered edge linking a composite scorer to one of its sub-scorers.""" + + __tablename__ = "ScorerIdentifierChildren" + __table_args__ = {"extend_existing": True} + + parent_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{ScorerIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + child_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{ScorerIdentifierEntry.__tablename__}.hash"), nullable=False + ) + + parent: Mapped["ScorerIdentifierEntry"] = relationship( + "ScorerIdentifierEntry", + foreign_keys=[parent_hash], + back_populates="sub_scorers", + ) + child: Mapped["ScorerIdentifierEntry"] = relationship( + "ScorerIdentifierEntry", + foreign_keys=[child_hash], + ) + + class ConversationEntry(Base): """ Conversation-scoped metadata, persisted once per ``conversation_id``. @@ -713,6 +778,9 @@ class ScoreEntry(Base): score_rationale = mapped_column(String, nullable=True) score_metadata: Mapped[dict[str, str | int | float]] = mapped_column(JSON) scorer_class_identifier: Mapped[dict[str, Any]] = mapped_column(JSON) + scorer_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{ScorerIdentifierEntry.__tablename__}.hash"), nullable=True + ) prompt_request_response_id = mapped_column(CustomUUID, ForeignKey(f"{PromptMemoryEntry.__tablename__}.id")) timestamp = mapped_column(UTCDateTime, nullable=False) task = mapped_column(String, nullable=True) # Deprecated: Use objective instead @@ -744,6 +812,7 @@ def __init__(self, *, entry: Score) -> None: ScorerEvaluationIdentifier(normalized_scorer).eval_hash ) self.scorer_class_identifier = normalized_scorer.model_dump() if normalized_scorer else {} + self.scorer_identifier_hash = normalized_scorer.hash if normalized_scorer else None self.prompt_request_response_id = entry.message_piece_id if entry.message_piece_id else None self.timestamp = entry.timestamp # Store in both columns for backward compatibility diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index 3a50b09ceb..85d934e923 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -755,6 +755,9 @@ def test_target_identifier_persists_inner_targets_and_edges(sqlite_instance: Mem class_module="pyrit.prompt_target", children={"targets": [inner_a, inner_b]}, ) + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id="conv-inner-a", target_identifier=inner_a) + ) sqlite_instance.add_conversation_to_memory( conversation=Conversation(conversation_id="conv-multi", target_identifier=multi) ) diff --git a/tests/unit/memory/memory_interface/test_interface_scores.py b/tests/unit/memory/memory_interface/test_interface_scores.py index 647fe75254..1cfde0f692 100644 --- a/tests/unit/memory/memory_interface/test_interface_scores.py +++ b/tests/unit/memory/memory_interface/test_interface_scores.py @@ -114,6 +114,73 @@ def test_add_score_get_score( assert db_score[0].message_piece_id == prompt_id +def test_scorer_identifier_persists_graph_and_dedupes( + sqlite_instance: MemoryInterface, + sample_conversation_entries: Sequence[PromptMemoryEntry], +): + from pyrit.memory.memory_models import ( + ScoreEntry, + ScorerIdentifierChildEntry, + ScorerIdentifierEntry, + TargetIdentifierEntry, + ) + + prompt_id = sample_conversation_entries[0].id + sqlite_instance._insert_entries(entries=sample_conversation_entries) + prompt_target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://example.test", "model_name": "gpt-test"}, + ) + sub_scorer = ComponentIdentifier( + class_name="SelfAskScorer", + class_module="pyrit.score", + params={"scorer_type": "float_scale"}, + children={"prompt_target": prompt_target}, + ) + composite = ComponentIdentifier( + class_name="CompositeScorer", + class_module="pyrit.score", + params={"scorer_type": "true_false", "score_aggregator": "AND_"}, + children={"sub_scorers": [sub_scorer]}, + ) + scores = [ + Score( + score_value="True", + score_type="true_false", + scorer_class_identifier=composite, + message_piece_id=prompt_id, + ), + Score( + score_value="True", + score_type="true_false", + scorer_class_identifier=composite, + message_piece_id=prompt_id, + ), + ] + + sqlite_instance.add_scores_to_memory(scores=scores) + + scorer_rows = sqlite_instance._query_entries(ScorerIdentifierEntry) + scorers_by_hash = {row.hash: row for row in scorer_rows} + assert set(scorers_by_hash) == {composite.hash, sub_scorer.hash} + assert scorers_by_hash[composite.hash].scorer_type == "true_false" + assert scorers_by_hash[composite.hash].score_aggregator == "AND_" + assert scorers_by_hash[composite.hash].prompt_target_hash is None + assert scorers_by_hash[sub_scorer.hash].prompt_target_hash == prompt_target.hash + assert ComponentIdentifier.model_validate(scorers_by_hash[composite.hash].identifier_json) == composite + + score_rows = sqlite_instance._query_entries(ScoreEntry) + assert {row.scorer_identifier_hash for row in score_rows} == {composite.hash} + + target_rows = sqlite_instance._query_entries(TargetIdentifierEntry) + assert [row.hash for row in target_rows] == [prompt_target.hash] + scorer_edges = sqlite_instance._query_entries(ScorerIdentifierChildEntry) + assert [(edge.parent_hash, edge.position, edge.child_hash) for edge in scorer_edges] == [ + (composite.hash, 0, sub_scorer.hash) + ] + + def test_get_prompt_scores_empty_prompt_ids_returns_empty(sqlite_instance: MemoryInterface): prompt_id = uuid4() piece = MessagePiece( diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 63a998778b..837cf9a45e 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -16,6 +16,7 @@ PromptMemoryEntry, ScenarioResultEntry, ScoreEntry, + ScorerIdentifierEntry, SeedEntry, UTCDateTime, _load_identifier, @@ -30,6 +31,7 @@ MessagePiece, ScenarioResult, Score, + ScorerIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, @@ -162,6 +164,35 @@ def test_load_identifier_injects_pyrit_version(): assert loaded.pyrit_version == "9.9.9" +def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): + leaf = ScorerIdentifier( + class_name="LeafScorer", + class_module="pyrit.score", + scorer_type="float_scale", + ) + nested = ScorerIdentifier( + class_name="NestedScorer", + class_module="pyrit.score", + scorer_type="true_false", + sub_scorers=[leaf], + ) + root = ScorerIdentifier( + class_name="RootScorer", + class_module="pyrit.score", + scorer_type="true_false", + sub_scorers=[nested, leaf], + ) + + entry = ScorerIdentifierEntry.from_domain_model(domain_model=root) + + assert [edge.position for edge in entry.sub_scorers] == [0, 1] + assert entry.sub_scorers[0].child.hash == nested.hash + assert entry.sub_scorers[1].child.hash == leaf.hash + nested_entry = entry.sub_scorers[0].child + assert len(nested_entry.sub_scorers) == 1 + assert nested_entry.sub_scorers[0].child is entry.sub_scorers[1].child + + # --------------------------------------------------------------------------- # ConversationMessageWithSimilarity # --------------------------------------------------------------------------- diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 015642ca89..91b0f6bc82 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import ast +import json import os import tempfile import uuid @@ -808,6 +809,76 @@ def test_conversations_migration_downgrade_restores_columns(): engine.dispose() +# ============================================================================= +# Backfill tests for scorer identifier persistence (a6c8e0f2b4d6) +# ============================================================================= + + +def test_scorer_identifier_migration_backfills_graph_and_score_link(): + """Existing scorer JSON is materialized and linked without changing its content identity.""" + from pyrit.models import ComponentIdentifier + + prompt_target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://example.test"}, + ) + sub_scorer = ComponentIdentifier( + class_name="SelfAskScorer", + class_module="pyrit.score", + children={"prompt_target": prompt_target}, + ) + composite = ComponentIdentifier( + class_name="CompositeScorer", + class_module="pyrit.score", + params={"scorer_type": "true_false", "score_aggregator": "AND_"}, + children={"sub_scorers": [sub_scorer]}, + ) + score_id = str(uuid.uuid4()) + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'scorer-backfill.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "e5f7a9c1b3d2") + connection.execute( + text( + 'INSERT INTO "ScoreEntries" ' + "(id, score_value, score_type, score_metadata, scorer_class_identifier, timestamp) " + "VALUES (:id, 'True', 'true_false', '{}', :identifier, '2026-07-13')" + ), + {"id": score_id, "identifier": json.dumps(composite.model_dump())}, + ) + + command.upgrade(config, "a6c8e0f2b4d6") + + score_hash = connection.execute( + text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), + {"id": score_id}, + ).scalar_one() + scorer_rows = connection.execute( + text( + 'SELECT hash, scorer_type, score_aggregator, prompt_target_hash FROM "ScorerIdentifiers"' + ) + ).fetchall() + target_hashes = connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars().all() + scorer_child = connection.execute( + text('SELECT parent_hash, position, child_hash FROM "ScorerIdentifierChildren"') + ).one() + + assert score_hash == composite.hash + assert {row[0] for row in scorer_rows} == {composite.hash, sub_scorer.hash} + root_row = next(row for row in scorer_rows if row[0] == composite.hash) + sub_scorer_row = next(row for row in scorer_rows if row[0] == sub_scorer.hash) + assert root_row[1:] == ("true_false", "AND_", None) + assert sub_scorer_row[3] == prompt_target.hash + assert target_hashes == [prompt_target.hash] + assert scorer_child == (composite.hash, 0, sub_scorer.hash) + finally: + engine.dispose() + + _STRING_TYPES_REQUIRING_LENGTH = {"String", "VARCHAR", "NVARCHAR", "Unicode"} From f27679ad29a8c3df94740e28fc1a2434fe59dcbd Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 14:52:34 -0700 Subject: [PATCH 04/24] add scenario identifier --- ...f1a3c5e7_add_scenario_identifiers_table.py | 220 ++++++++++++++++++ pyrit/memory/memory_interface.py | 44 +++- pyrit/memory/memory_models.py | 45 ++++ .../test_interface_scenario_results.py | 49 ++++ tests/unit/memory/test_migration.py | 91 ++++++++ 5 files changed, 446 insertions(+), 3 deletions(-) create mode 100644 pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py new file mode 100644 index 0000000000..772c1ffe6c --- /dev/null +++ b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py @@ -0,0 +1,220 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist scenario identifiers as content-addressed rows. + +Revision ID: b7d9f1a3c5e7 +Revises: a6c8e0f2b4d6 +Create Date: 2026-07-13 13:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +revision: str = "b7d9f1a3c5e7" +down_revision: str | None = "a6c8e0f2b4d6" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "ScenarioIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("version", sa.Integer(), nullable=True), + sa.Column("techniques", sa.JSON(), nullable=True), + sa.Column("datasets", sa.JSON(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["objective_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_scenario_identifiers_objective_target_hash", + ), + sa.ForeignKeyConstraint( + ["objective_scorer_hash"], + ["ScorerIdentifiers.hash"], + name="fk_scenario_identifiers_objective_scorer_hash", + ), + ) + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.add_column(sa.Column("scenario_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_scenario_result_entries_scenario_identifier_hash", + "ScenarioIdentifiers", + ["scenario_identifier_hash"], + ["hash"], + ) + + _backfill_scenario_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.drop_column("scenario_identifier_hash") + op.drop_table("ScenarioIdentifiers") + + +def _backfill_scenario_identifiers() -> None: + """Backfill scenario rows and result foreign keys from the retained JSON column.""" + from pyrit.models import ScenarioIdentifier, ScorerIdentifier, TargetIdentifier + + bind = op.get_bind() + result_rows = bind.execute( + sa.text( + 'SELECT id, scenario_identifier FROM "ScenarioResultEntries" ' + "WHERE scenario_identifier IS NOT NULL" + ) + ).fetchall() + existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScenarioIdentifiers"')).fetchall()} + target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} + scorer_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScorerIdentifiers"')).fetchall()} + target_insert = sa.text( + 'INSERT INTO "TargetIdentifiers" ' + "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " + "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " + ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " + ":supported_auth_modes, :pyrit_version)" + ) + target_edge_insert = sa.text( + 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + scorer_insert = sa.text( + 'INSERT INTO "ScorerIdentifiers" ' + "(hash, class_name, class_module, identifier_json, scorer_type, score_aggregator, " + "prompt_target_hash, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :scorer_type, " + ":score_aggregator, :prompt_target_hash, :pyrit_version)" + ) + scorer_edge_insert = sa.text( + 'INSERT INTO "ScorerIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + insert_stmt = sa.text( + 'INSERT INTO "ScenarioIdentifiers" ' + "(hash, class_name, class_module, identifier_json, version, techniques, datasets, " + "objective_target_hash, objective_scorer_hash, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :version, :techniques, :datasets, " + ":objective_target_hash, :objective_scorer_hash, :pyrit_version)" + ) + update_stmt = sa.text( + 'UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id' + ) + + def _insert_target(identifier: TargetIdentifier) -> None: + if identifier.hash in target_hashes: + return + for child in identifier.targets: + _insert_target(child) + bind.execute( + target_insert, + { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "endpoint": identifier.endpoint, + "model_name": identifier.model_name, + "underlying_model_name": identifier.underlying_model_name, + "temperature": identifier.temperature, + "top_p": identifier.top_p, + "max_requests_per_minute": identifier.max_requests_per_minute, + "supported_auth_modes": ( + json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None + ), + "pyrit_version": identifier.pyrit_version, + }, + ) + target_hashes.add(identifier.hash) + for position, child in enumerate(identifier.targets): + bind.execute( + target_edge_insert, + {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, + ) + + def _insert_scorer(identifier: ScorerIdentifier) -> None: + if identifier.hash in scorer_hashes: + return + if identifier.prompt_target is not None: + _insert_target(identifier.prompt_target) + for child in identifier.sub_scorers: + _insert_scorer(child) + bind.execute( + scorer_insert, + { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "scorer_type": identifier.scorer_type, + "score_aggregator": identifier.score_aggregator, + "prompt_target_hash": identifier.prompt_target.hash if identifier.prompt_target is not None else None, + "pyrit_version": identifier.pyrit_version, + }, + ) + scorer_hashes.add(identifier.hash) + for position, child in enumerate(identifier.sub_scorers): + bind.execute( + scorer_edge_insert, + {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, + ) + + skipped = 0 + for result_id, raw_scenario in result_rows: + stored = json.loads(raw_scenario) if isinstance(raw_scenario, str) else raw_scenario + if not stored: + continue + try: + identifier = ScenarioIdentifier.model_validate(stored) + if identifier.objective_target is not None: + _insert_target(identifier.objective_target) + if identifier.objective_scorer is not None: + _insert_scorer(identifier.objective_scorer) + if identifier.hash not in existing_hashes: + bind.execute( + insert_stmt, + { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "version": identifier.version, + "techniques": json.dumps(identifier.techniques) if identifier.techniques is not None else None, + "datasets": json.dumps(identifier.datasets) if identifier.datasets is not None else None, + "objective_target_hash": ( + identifier.objective_target.hash if identifier.objective_target is not None else None + ), + "objective_scorer_hash": ( + identifier.objective_scorer.hash if identifier.objective_scorer is not None else None + ), + "pyrit_version": identifier.pyrit_version, + }, + ) + existing_hashes.add(identifier.hash) + bind.execute(update_stmt, {"hash": identifier.hash, "id": result_id}) + except Exception: + skipped += 1 + logger.warning( + f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", + exc_info=True, + ) + + if skipped: + logger.warning(f"ScenarioIdentifiers backfill skipped {skipped} scenario result row(s)") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index b28dbd6ce8..ea0996fc8e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -26,6 +26,7 @@ ConversationEntry, EmbeddingDataEntry, PromptMemoryEntry, + ScenarioIdentifierEntry, ScenarioResultEntry, ScoreEntry, ScorerIdentifierEntry, @@ -46,6 +47,7 @@ IdentifierType, Message, MessagePiece, + ScenarioIdentifier, ScenarioResult, Score, ScorerIdentifier, @@ -2213,10 +2215,46 @@ def add_scenario_results_to_memory(self, *, scenario_results: Sequence[ScenarioR Args: scenario_results: Sequence of ScenarioResult objects to store in the database. + + Raises: + SQLAlchemyError: If a scenario result or identifier graph cannot be persisted. """ - self._insert_entries( - entries=[ScenarioResultEntry(entry=scenario_result) for scenario_result in scenario_results] - ) + entries = [ScenarioResultEntry(entry=scenario_result) for scenario_result in scenario_results] + with closing(self.get_session()) as session: + try: + for scenario_result in scenario_results: + self._persist_scenario_identifier( + session=session, + scenario_identifier=scenario_result.scenario_identifier, + ) + session.add_all(entries) + session.commit() + except SQLAlchemyError: + session.rollback() + raise + + @staticmethod + def _persist_scenario_identifier(*, session: Any, scenario_identifier: ScenarioIdentifier) -> None: + """Persist a scenario identifier and its target and scorer dependencies.""" + if session.get(ScenarioIdentifierEntry, scenario_identifier.hash) is not None: + return + if scenario_identifier.objective_target is not None: + MemoryInterface._persist_target_identifier( + session=session, + target_identifier=scenario_identifier.objective_target, + ) + if scenario_identifier.objective_scorer is not None: + MemoryInterface._persist_scorer_identifier( + session=session, + scorer_identifier=scenario_identifier.objective_scorer, + ) + try: + with session.begin_nested(): + entry = ScenarioIdentifierEntry.from_domain_model(scenario_identifier) + session.merge(entry) + session.flush() + except IntegrityError: + pass def update_scenario_run_state( self, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index ce81d3d75e..ad3590b84c 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -678,6 +678,43 @@ class ScorerIdentifierChildEntry(Base): ) +class ScenarioIdentifierEntry(ComponentIdentifierEntry[ScenarioIdentifier]): + """Content-addressed store of ``ScenarioIdentifier`` projections.""" + + __tablename__ = "ScenarioIdentifiers" + __table_args__ = {"extend_existing": True} + + version: Mapped[int | None] = mapped_column(INTEGER, nullable=True) + techniques: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + datasets: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + objective_target_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) + objective_scorer_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{ScorerIdentifierEntry.__tablename__}.hash"), nullable=True + ) + + objective_target: Mapped["TargetIdentifierEntry | None"] = relationship( + "TargetIdentifierEntry", + foreign_keys=[objective_target_hash], + ) + objective_scorer: Mapped["ScorerIdentifierEntry | None"] = relationship( + "ScorerIdentifierEntry", + foreign_keys=[objective_scorer_hash], + ) + + @classmethod + def _from_domain_model_shallow(cls, *, domain_model: ScenarioIdentifier) -> Self: + entry = super()._from_domain_model_shallow(domain_model=domain_model) + entry.objective_target_hash = ( + domain_model.objective_target.hash if domain_model.objective_target is not None else None + ) + entry.objective_scorer_hash = ( + domain_model.objective_scorer.hash if domain_model.objective_scorer is not None else None + ) + return entry + + class ConversationEntry(Base): """ Conversation-scoped metadata, persisted once per ``conversation_id``. @@ -1474,6 +1511,13 @@ class ScenarioResultEntry(Base): #: Canonical scenario identity (class name, version, techniques, datasets, #: resolved params, objective target / scorer children) with its eval hash. scenario_identifier: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) + scenario_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{ScenarioIdentifierEntry.__tablename__}.hash"), nullable=True + ) + scenario_identifier_entry: Mapped["ScenarioIdentifierEntry | None"] = relationship( + "ScenarioIdentifierEntry", + foreign_keys=[scenario_identifier_hash], + ) objective_target_identifier: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) objective_scorer_identifier: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) scenario_run_state: Mapped[str] = mapped_column(String, nullable=False, default="CREATED") @@ -1516,6 +1560,7 @@ def __init__(self, *, entry: ScenarioResult) -> None: ScenarioEvaluationIdentifier(entry.scenario_identifier).eval_hash ) self.scenario_identifier = scenario_identifier.model_dump() + self.scenario_identifier_hash = scenario_identifier.hash # Convert ComponentIdentifier to dict for JSON storage target_identifier = entry.objective_target_identifier diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index d1dae1ac38..8e3980fabc 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -7,12 +7,20 @@ from unit.mocks import get_mock_scorer_identifier, make_scenario_result from pyrit.memory import MemoryInterface +from pyrit.memory.memory_models import ( + ScenarioIdentifierEntry, + ScenarioResultEntry, + ScorerIdentifierEntry, + TargetIdentifierEntry, +) from pyrit.models import ( AttackOutcome, AttackResult, ComponentIdentifier, IdentifierFilter, IdentifierType, + ScorerIdentifier, + TargetIdentifier, ) @@ -89,6 +97,47 @@ def test_add_and_retrieve_scenario_results(sqlite_instance: MemoryInterface, sam assert scenario_names == {"Scenario 1", "Scenario 2"} +def test_add_scenario_results_persists_identifier_graph(sqlite_instance: MemoryInterface): + target = TargetIdentifier( + class_name="TestTarget", + class_module="tests.unit.memory", + model_name="test-model", + ) + scorer = ScorerIdentifier( + class_name="TestScorer", + class_module="tests.unit.memory", + scorer_type="true_false", + prompt_target=target, + ) + results = [ + make_scenario_result( + scenario_name="PersistedScenario", + scenario_version=2, + techniques=["TechniqueA"], + datasets=["DatasetA"], + objective_target_identifier=target, + objective_scorer_identifier=scorer, + attack_results={}, + ) + for _ in range(2) + ] + + sqlite_instance.add_scenario_results_to_memory(scenario_results=results) + + scenario_rows = sqlite_instance._query_entries(ScenarioIdentifierEntry) + result_rows = sqlite_instance._query_entries(ScenarioResultEntry) + assert len(scenario_rows) == 1 + assert scenario_rows[0].hash == results[0].scenario_identifier.hash + assert scenario_rows[0].version == 2 + assert scenario_rows[0].techniques == ["TechniqueA"] + assert scenario_rows[0].datasets == ["DatasetA"] + assert scenario_rows[0].objective_target_hash == target.hash + assert scenario_rows[0].objective_scorer_hash == scorer.hash + assert {row.scenario_identifier_hash for row in result_rows} == {scenario_rows[0].hash} + assert len(sqlite_instance._query_entries(TargetIdentifierEntry)) == 1 + assert len(sqlite_instance._query_entries(ScorerIdentifierEntry)) == 1 + + def test_filter_by_name(sqlite_instance: MemoryInterface, sample_attack_results): """Test retrieving scenario results filtered by name.""" # Create and add scenario results diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 91b0f6bc82..e168344bea 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -879,6 +879,97 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): engine.dispose() +# ============================================================================= +# Backfill tests for scenario identifier persistence (b7d9f1a3c5e7) +# ============================================================================= + + +def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): + """Scenario-only target and scorer graphs are materialized before the scenario row.""" + from pyrit.models import ScenarioIdentifier, ScorerIdentifier, TargetIdentifier + + target = TargetIdentifier( + class_name="ObjectiveTarget", + class_module="pyrit.prompt_target", + model_name="objective-model", + ) + scorer_target = TargetIdentifier( + class_name="ScorerTarget", + class_module="pyrit.prompt_target", + model_name="scorer-model", + ) + scorer = ScorerIdentifier( + class_name="SelfAskScorer", + class_module="pyrit.score", + scorer_type="true_false", + prompt_target=scorer_target, + ) + scenario = ScenarioIdentifier( + class_name="TestScenario", + class_module="pyrit.scenario", + version=3, + techniques=["TechniqueA"], + datasets=["DatasetA"], + objective_target=target, + objective_scorer=scorer, + ) + result_id = str(uuid.uuid4()) + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'scenario-backfill.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "a6c8e0f2b4d6") + connection.execute( + text( + 'INSERT INTO "ScenarioResultEntries" ' + "(id, scenario_name, scenario_version, pyrit_version, scenario_identifier, " + "objective_target_identifier, objective_scorer_identifier, scenario_run_state, " + "attack_results_json, number_tries, completion_time, timestamp) " + "VALUES (:id, 'TestScenario', 3, '0.10.0', :scenario, :target, :scorer, " + "'COMPLETED', '{}', 1, '2026-07-13', '2026-07-13')" + ), + { + "id": result_id, + "scenario": json.dumps(scenario.model_dump()), + "target": json.dumps(target.model_dump()), + "scorer": json.dumps(scorer.model_dump()), + }, + ) + + command.upgrade(config, "b7d9f1a3c5e7") + + result_hash = connection.execute( + text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() + scenario_row = connection.execute( + text( + 'SELECT hash, version, techniques, datasets, objective_target_hash, objective_scorer_hash ' + 'FROM "ScenarioIdentifiers"' + ) + ).one() + target_hashes = set(connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars()) + scorer_row = connection.execute( + text('SELECT hash, prompt_target_hash FROM "ScorerIdentifiers"') + ).one() + + assert result_hash == scenario.hash + assert scenario_row == ( + scenario.hash, + 3, + json.dumps(["TechniqueA"]), + json.dumps(["DatasetA"]), + target.hash, + scorer.hash, + ) + assert target_hashes == {target.hash, scorer_target.hash} + assert scorer_row == (scorer.hash, scorer_target.hash) + finally: + engine.dispose() + + _STRING_TYPES_REQUIRING_LENGTH = {"String", "VARCHAR", "NVARCHAR", "Unicode"} From 1b098c1d56d84b82bef3a5cbc3ff5a0be5c6253f Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 15:13:49 -0700 Subject: [PATCH 05/24] refactor --- pyrit/memory/memory_interface.py | 92 +++++++++++-------------- pyrit/memory/memory_models.py | 32 ++++----- tests/unit/memory/test_memory_models.py | 31 +++++++++ 3 files changed, 87 insertions(+), 68 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index ea0996fc8e..fc115a3b86 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -7,7 +7,7 @@ import re import uuid import weakref -from collections.abc import MutableSequence, Sequence +from collections.abc import Iterator, MutableSequence, Sequence from contextlib import closing from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar @@ -23,6 +23,7 @@ from pyrit.memory.memory_models import ( AttackResultEntry, Base, + ComponentIdentifierEntry, ConversationEntry, EmbeddingDataEntry, PromptMemoryEntry, @@ -41,6 +42,7 @@ ) from pyrit.models import ( AttackResult, + ComponentIdentifier, Conversation, ConversationStats, IdentifierFilter, @@ -489,8 +491,8 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}") raise - @staticmethod - def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentifier) -> None: + @classmethod + def _persist_target_identifier(cls, *, session: Any, target_identifier: TargetIdentifier) -> None: """ Persist ``target_identifier`` and its inner targets as content-addressed rows. @@ -511,18 +513,43 @@ def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentif session (Any): The active SQLAlchemy session (the caller's transaction). target_identifier (TargetIdentifier): The target identifier to persist. """ - if session.get(TargetIdentifierEntry, target_identifier.hash) is not None: + cls._persist_identifier(session=session, identifier=target_identifier) + + @classmethod + def _persist_identifier(cls, *, session: Any, identifier: ComponentIdentifier) -> None: + entry_type = cls._get_identifier_entry_type(identifier) + if session.get(entry_type, identifier.hash) is not None: return + + for dependency in cls._iter_identifier_dependencies(identifier): + cls._persist_identifier(session=session, identifier=dependency) + try: with session.begin_nested(): - entry = TargetIdentifierEntry.from_domain_model(target_identifier) - session.merge(entry) + session.merge(entry_type.from_domain_model(identifier)) session.flush() except IntegrityError: - # A racing writer inserted the same content-addressed row; immutable - # rows make this a safe no-op. pass + @staticmethod + def _get_identifier_entry_type(identifier: ComponentIdentifier) -> type[ComponentIdentifierEntry[Any]]: + if isinstance(identifier, TargetIdentifier): + return TargetIdentifierEntry + if isinstance(identifier, ScorerIdentifier): + return ScorerIdentifierEntry + if isinstance(identifier, ScenarioIdentifier): + return ScenarioIdentifierEntry + raise TypeError(f"Identifier type {type(identifier).__name__} does not have a persistence entry.") + + @staticmethod + def _iter_identifier_dependencies(identifier: ComponentIdentifier) -> Iterator[ComponentIdentifier]: + for field_name in identifier.promoted_child_field_names(): + child = getattr(identifier, field_name) + if isinstance(child, ComponentIdentifier): + yield child + elif isinstance(child, list): + yield from (item for item in child if isinstance(item, ComponentIdentifier)) + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ @@ -881,29 +908,10 @@ def add_scores_to_memory(self, *, scores: Sequence[Score]) -> None: logger.exception(f"Error inserting scores: {e}") raise - @staticmethod - def _persist_scorer_identifier(*, session: Any, scorer_identifier: ScorerIdentifier) -> None: + @classmethod + def _persist_scorer_identifier(cls, *, session: Any, scorer_identifier: ScorerIdentifier) -> None: """Persist a complete scorer graph and its target dependencies.""" - if session.get(ScorerIdentifierEntry, scorer_identifier.hash) is not None: - return - - pending_scorers = [scorer_identifier] - while pending_scorers: - pending_scorer = pending_scorers.pop() - if pending_scorer.prompt_target is not None: - MemoryInterface._persist_target_identifier( - session=session, - target_identifier=pending_scorer.prompt_target, - ) - pending_scorers.extend(pending_scorer.sub_scorers) - - try: - with session.begin_nested(): - entry = ScorerIdentifierEntry.from_domain_model(scorer_identifier) - session.merge(entry) - session.flush() - except IntegrityError: - pass + cls._persist_identifier(session=session, identifier=scorer_identifier) def get_scores( self, @@ -2233,28 +2241,10 @@ def add_scenario_results_to_memory(self, *, scenario_results: Sequence[ScenarioR session.rollback() raise - @staticmethod - def _persist_scenario_identifier(*, session: Any, scenario_identifier: ScenarioIdentifier) -> None: + @classmethod + def _persist_scenario_identifier(cls, *, session: Any, scenario_identifier: ScenarioIdentifier) -> None: """Persist a scenario identifier and its target and scorer dependencies.""" - if session.get(ScenarioIdentifierEntry, scenario_identifier.hash) is not None: - return - if scenario_identifier.objective_target is not None: - MemoryInterface._persist_target_identifier( - session=session, - target_identifier=scenario_identifier.objective_target, - ) - if scenario_identifier.objective_scorer is not None: - MemoryInterface._persist_scorer_identifier( - session=session, - scorer_identifier=scenario_identifier.objective_scorer, - ) - try: - with session.begin_nested(): - entry = ScenarioIdentifierEntry.from_domain_model(scenario_identifier) - session.merge(entry) - session.flush() - except IntegrityError: - pass + cls._persist_identifier(session=session, identifier=scenario_identifier) def update_scenario_run_state( self, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index ad3590b84c..60f927b9d1 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -449,6 +449,8 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): #: Optional per-child-field wiring for materialized edge relationships. #: Empty by default so identifier rows only persist their own projection. CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = {} + #: Mapping from singular promoted child fields to their foreign-key columns. + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = {} #: Content-addressed identity — the same value as ``ComponentIdentifier.hash``. #: SHA256 hex digest is 64 chars; bounded for SQL Server key/index compatibility. @@ -496,8 +498,15 @@ def _from_domain_model_shallow(cls, *, domain_model: T) -> Self: ) for name, value in domain_model.promoted_scalar_values().items(): setattr(entry, name, value) # each promoted scalar → its mapped column + cls._populate_child_hashes(entry=entry, domain_model=domain_model) return entry + @classmethod + def _populate_child_hashes(cls, *, entry: Self, domain_model: T) -> None: + for child_field, hash_column in cls.CHILD_HASH_COLUMNS.items(): + child = getattr(domain_model, child_field) + setattr(entry, hash_column, child.hash if child is not None else None) + @classmethod def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: dict[str, Self]) -> None: for field_name in domain_model.promoted_child_field_names(): @@ -627,6 +636,7 @@ class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): edge_position_attr="position", ) } + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = {"prompt_target": "prompt_target_hash"} scorer_type: Mapped[str | None] = mapped_column(String, nullable=True) score_aggregator: Mapped[str | None] = mapped_column(String, nullable=True) @@ -646,12 +656,6 @@ class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): back_populates="parent", ) - @classmethod - def _from_domain_model_shallow(cls, *, domain_model: ScorerIdentifier) -> Self: - entry = super()._from_domain_model_shallow(domain_model=domain_model) - entry.prompt_target_hash = domain_model.prompt_target.hash if domain_model.prompt_target is not None else None - return entry - class ScorerIdentifierChildEntry(Base): """Ordered edge linking a composite scorer to one of its sub-scorers.""" @@ -684,6 +688,11 @@ class ScenarioIdentifierEntry(ComponentIdentifierEntry[ScenarioIdentifier]): __tablename__ = "ScenarioIdentifiers" __table_args__ = {"extend_existing": True} + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = { + "objective_target": "objective_target_hash", + "objective_scorer": "objective_scorer_hash", + } + version: Mapped[int | None] = mapped_column(INTEGER, nullable=True) techniques: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) datasets: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) @@ -703,17 +712,6 @@ class ScenarioIdentifierEntry(ComponentIdentifierEntry[ScenarioIdentifier]): foreign_keys=[objective_scorer_hash], ) - @classmethod - def _from_domain_model_shallow(cls, *, domain_model: ScenarioIdentifier) -> Self: - entry = super()._from_domain_model_shallow(domain_model=domain_model) - entry.objective_target_hash = ( - domain_model.objective_target.hash if domain_model.objective_target is not None else None - ) - entry.objective_scorer_hash = ( - domain_model.objective_scorer.hash if domain_model.objective_scorer is not None else None - ) - return entry - class ConversationEntry(Base): """ diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 837cf9a45e..05bc00a65e 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -2,7 +2,9 @@ # Licensed under the MIT license. import uuid +from collections.abc import Sequence from datetime import datetime, timezone +from typing import Any, get_origin from unittest.mock import MagicMock import pytest @@ -10,14 +12,17 @@ from pyrit.memory.memory_models import ( AttackResultEntry, + ComponentIdentifierEntry, ConversationMessageWithSimilarity, EmbeddingDataEntry, EmbeddingMessageWithSimilarity, PromptMemoryEntry, + ScenarioIdentifierEntry, ScenarioResultEntry, ScoreEntry, ScorerIdentifierEntry, SeedEntry, + TargetIdentifierEntry, UTCDateTime, _load_identifier, ) @@ -29,12 +34,14 @@ ConversationReference, ConversationType, MessagePiece, + ScenarioIdentifier, ScenarioResult, Score, ScorerIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, + TargetIdentifier, ) from unit.mocks import make_scenario_result @@ -193,6 +200,30 @@ def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): assert nested_entry.sub_scorers[0].child is entry.sub_scorers[1].child +@pytest.mark.parametrize( + ("identifier_type", "entry_type"), + [ + (TargetIdentifier, TargetIdentifierEntry), + (ScorerIdentifier, ScorerIdentifierEntry), + (ScenarioIdentifier, ScenarioIdentifierEntry), + ], +) +def test_identifier_entry_maps_promoted_children_by_cardinality( + identifier_type: type[ComponentIdentifier], + entry_type: type[ComponentIdentifierEntry[Any]], +) -> None: + promoted_children = set(identifier_type.promoted_child_field_names()) + collection_children = { + field_name + for field_name in promoted_children + if get_origin(identifier_type.model_fields[field_name].annotation) in (list, Sequence) + } + singular_children = promoted_children - collection_children + + assert set(entry_type.CHILD_RELATIONSHIP_SPECS) == collection_children + assert set(entry_type.CHILD_HASH_COLUMNS) == singular_children + + # --------------------------------------------------------------------------- # ConversationMessageWithSimilarity # --------------------------------------------------------------------------- From 1838ca01841d424f5188547232da848e31f3bde4 Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 15:15:16 -0700 Subject: [PATCH 06/24] ruff --- .../a6c8e0f2b4d6_add_scorer_identifiers_table.py | 5 +---- .../b7d9f1a3c5e7_add_scenario_identifiers_table.py | 9 ++------- tests/unit/memory/test_migration.py | 10 +++------- 3 files changed, 6 insertions(+), 18 deletions(-) diff --git a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py index 184185e2f2..669a64716c 100644 --- a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py +++ b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py @@ -81,10 +81,7 @@ def _backfill_scorer_identifiers() -> None: bind = op.get_bind() score_rows = bind.execute( - sa.text( - 'SELECT id, scorer_class_identifier FROM "ScoreEntries" ' - "WHERE scorer_class_identifier IS NOT NULL" - ) + sa.text('SELECT id, scorer_class_identifier FROM "ScoreEntries" WHERE scorer_class_identifier IS NOT NULL') ).fetchall() scorer_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScorerIdentifiers"')).fetchall()} target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} diff --git a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py index 772c1ffe6c..859ef211c0 100644 --- a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py +++ b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py @@ -76,10 +76,7 @@ def _backfill_scenario_identifiers() -> None: bind = op.get_bind() result_rows = bind.execute( - sa.text( - 'SELECT id, scenario_identifier FROM "ScenarioResultEntries" ' - "WHERE scenario_identifier IS NOT NULL" - ) + sa.text('SELECT id, scenario_identifier FROM "ScenarioResultEntries" WHERE scenario_identifier IS NOT NULL') ).fetchall() existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScenarioIdentifiers"')).fetchall()} target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} @@ -114,9 +111,7 @@ def _backfill_scenario_identifiers() -> None: "VALUES (:hash, :class_name, :class_module, :identifier_json, :version, :techniques, :datasets, " ":objective_target_hash, :objective_scorer_hash, :pyrit_version)" ) - update_stmt = sa.text( - 'UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id' - ) + update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') def _insert_target(identifier: TargetIdentifier) -> None: if identifier.hash in target_hashes: diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index e168344bea..ef1c04e59b 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -858,9 +858,7 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): {"id": score_id}, ).scalar_one() scorer_rows = connection.execute( - text( - 'SELECT hash, scorer_type, score_aggregator, prompt_target_hash FROM "ScorerIdentifiers"' - ) + text('SELECT hash, scorer_type, score_aggregator, prompt_target_hash FROM "ScorerIdentifiers"') ).fetchall() target_hashes = connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars().all() scorer_child = connection.execute( @@ -946,14 +944,12 @@ def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): ).scalar_one() scenario_row = connection.execute( text( - 'SELECT hash, version, techniques, datasets, objective_target_hash, objective_scorer_hash ' + "SELECT hash, version, techniques, datasets, objective_target_hash, objective_scorer_hash " 'FROM "ScenarioIdentifiers"' ) ).one() target_hashes = set(connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars()) - scorer_row = connection.execute( - text('SELECT hash, prompt_target_hash FROM "ScorerIdentifiers"') - ).one() + scorer_row = connection.execute(text('SELECT hash, prompt_target_hash FROM "ScorerIdentifiers"')).one() assert result_hash == scenario.hash assert scenario_row == ( From 82642e5f0b2acc3e211ba7bfa0a14118c6354ce2 Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 16:48:45 -0700 Subject: [PATCH 07/24] persist converter identifiers --- ...3a5b7d9_add_converter_identifiers_table.py | 214 ++++++++++++++++++ pyrit/memory/azure_sql_memory.py | 17 +- pyrit/memory/memory_interface.py | 26 +++ pyrit/memory/memory_models.py | 62 +++++ pyrit/memory/sqlite_memory.py | 12 +- .../test_interface_prompts.py | 52 +++++ tests/unit/memory/test_memory_models.py | 3 + tests/unit/memory/test_migration.py | 74 ++++++ 8 files changed, 433 insertions(+), 27 deletions(-) create mode 100644 pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py b/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py new file mode 100644 index 0000000000..de0b529631 --- /dev/null +++ b/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py @@ -0,0 +1,214 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist converter identifiers as content-addressed rows. + +Revision ID: c8e1f3a5b7d9 +Revises: b7d9f1a3c5e7 +Create Date: 2026-07-13 15:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +import uuid +from collections.abc import Sequence # noqa: TC003 +from typing import TYPE_CHECKING, Any + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.sqlite import CHAR +from sqlalchemy.types import TypeDecorator, Uuid + +if TYPE_CHECKING: + from sqlalchemy.engine import Dialect + + +class _CustomUUID(TypeDecorator[uuid.UUID]): + """Frozen UUID type matching ``PromptMemoryEntries.id`` across dialects.""" + + impl = CHAR + cache_ok = True + + def load_dialect_impl(self, dialect: Dialect) -> Any: + if dialect.name == "sqlite": + return dialect.type_descriptor(CHAR(36)) + return dialect.type_descriptor(Uuid()) + + def process_bind_param(self, value: Any, dialect: Any) -> str | None: + return str(value) if value is not None else None + + def process_result_value(self, value: Any, dialect: Any) -> uuid.UUID | None: + if value is None: + return None + if isinstance(value, uuid.UUID): + return value + return uuid.UUID(value) + + +revision: str = "c8e1f3a5b7d9" +down_revision: str | None = "b7d9f1a3c5e7" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "ConverterIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("supported_input_types", sa.JSON(), nullable=True), + sa.Column("supported_output_types", sa.JSON(), nullable=True), + sa.Column("converter_target_hash", sa.String(64), nullable=True), + sa.Column("sub_converter_hash", sa.String(64), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["converter_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_converter_identifiers_converter_target_hash", + ), + sa.ForeignKeyConstraint( + ["sub_converter_hash"], + ["ConverterIdentifiers.hash"], + name="fk_converter_identifiers_sub_converter_hash", + ), + ) + op.create_table( + "PromptConverterIdentifiers", + sa.Column("prompt_memory_entry_id", _CustomUUID(), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("converter_identifier_hash", sa.String(64), nullable=False), + sa.ForeignKeyConstraint( + ["prompt_memory_entry_id"], + ["PromptMemoryEntries.id"], + name="fk_prompt_converter_identifiers_prompt_memory_entry_id", + ), + sa.ForeignKeyConstraint( + ["converter_identifier_hash"], + ["ConverterIdentifiers.hash"], + name="fk_prompt_converter_identifiers_converter_identifier_hash", + ), + sa.PrimaryKeyConstraint("prompt_memory_entry_id", "position"), + ) + _backfill_converter_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + op.drop_table("PromptConverterIdentifiers") + op.drop_table("ConverterIdentifiers") + + +def _backfill_converter_identifiers() -> None: + """Materialize converter graphs and prompt associations from retained JSON.""" + from pyrit.models import ComponentIdentifier, ConverterIdentifier, TargetIdentifier + + bind = op.get_bind() + prompt_rows = bind.execute( + sa.text( + 'SELECT id, converter_identifiers, pyrit_version FROM "PromptMemoryEntries" ' + "WHERE converter_identifiers IS NOT NULL" + ) + ).fetchall() + converter_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ConverterIdentifiers"'))} + target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"'))} + + target_insert = sa.text( + 'INSERT INTO "TargetIdentifiers" ' + "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " + "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " + ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " + ":supported_auth_modes, :pyrit_version)" + ) + target_edge_insert = sa.text( + 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + converter_insert = sa.text( + 'INSERT INTO "ConverterIdentifiers" ' + "(hash, class_name, class_module, identifier_json, supported_input_types, supported_output_types, " + "converter_target_hash, sub_converter_hash, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :supported_input_types, " + ":supported_output_types, :converter_target_hash, :sub_converter_hash, :pyrit_version)" + ) + link_insert = sa.text( + 'INSERT INTO "PromptConverterIdentifiers" ' + "(prompt_memory_entry_id, position, converter_identifier_hash) " + "VALUES (:prompt_memory_entry_id, :position, :converter_identifier_hash)" + ) + + def _insert_target(identifier: TargetIdentifier) -> None: + if identifier.hash in target_hashes: + return + for child in identifier.targets: + _insert_target(child) + bind.execute(target_insert, _identifier_values(identifier)) + for position, child in enumerate(identifier.targets): + bind.execute( + target_edge_insert, + {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, + ) + target_hashes.add(identifier.hash) + + def _insert_converter(identifier: ConverterIdentifier) -> None: + if identifier.hash in converter_hashes: + return + if identifier.converter_target is not None: + _insert_target(identifier.converter_target) + if identifier.sub_converter is not None: + _insert_converter(identifier.sub_converter) + bind.execute(converter_insert, _identifier_values(identifier)) + converter_hashes.add(identifier.hash) + + skipped = 0 + for prompt_id, stored_identifiers, pyrit_version in prompt_rows: + try: + values = json.loads(stored_identifiers) if isinstance(stored_identifiers, str) else stored_identifiers + for position, stored_identifier in enumerate(values): + stored_identifier["pyrit_version"] = pyrit_version + identifier = ConverterIdentifier.from_component_identifier( + ComponentIdentifier.model_validate(stored_identifier) + ) + _insert_converter(identifier) + bind.execute( + link_insert, + { + "prompt_memory_entry_id": prompt_id, + "position": position, + "converter_identifier_hash": identifier.hash, + }, + ) + except (TypeError, ValueError, KeyError): + skipped += 1 + logger.warning(f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}") + if skipped: + logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") + + +def _identifier_values(identifier: Any) -> dict[str, Any]: + """Return common and promoted values for a normalized identifier insert.""" + promoted_values = { + name: json.dumps(value) if isinstance(value, (list, dict)) else value + for name, value in identifier.promoted_scalar_values().items() + } + return { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump()), + "pyrit_version": identifier.pyrit_version, + **promoted_values, + **{ + f"{field_name}_hash": child.hash if child is not None else None + for field_name in identifier.promoted_child_field_names() + if not isinstance((child := getattr(identifier, field_name)), list) + }, + } diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 998ae0480a..cbd0fb187d 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -27,7 +27,7 @@ PromptMemoryEntry, ) from pyrit.memory.storage import AzureBlobStorageIO -from pyrit.models import ConversationStats, MessagePiece +from pyrit.models import ConversationStats if TYPE_CHECKING: from azure.core.credentials import AccessToken @@ -685,21 +685,6 @@ def _get_scenario_result_label_condition(self, *, labels: dict[str, str]) -> Any conditions.append(condition) return and_(*conditions) - def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece]) -> None: - """ - Persist already-validated message pieces to the Azure SQL store. - - ``not_in_memory`` pieces are ephemeral -- typically synthesized inside a - scorer to score arbitrary content that never came through a real - PromptTarget. They are filtered out upstream in - ``add_message_pieces_to_memory`` before this method is called. - - Args: - message_pieces (Sequence[MessagePiece]): Persistable pieces (filtered and - validated by ``add_message_pieces_to_memory``). - """ - self._insert_entries(entries=[PromptMemoryEntry(entry=piece) for piece in message_pieces]) - def dispose_engine(self) -> None: """ Dispose the engine and clean up resources. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index fc115a3b86..8c6eb0b152 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -25,7 +25,9 @@ Base, ComponentIdentifierEntry, ConversationEntry, + ConverterIdentifierEntry, EmbeddingDataEntry, + PromptConverterIdentifierEntry, PromptMemoryEntry, ScenarioIdentifierEntry, ScenarioResultEntry, @@ -45,6 +47,7 @@ ComponentIdentifier, Conversation, ConversationStats, + ConverterIdentifier, IdentifierFilter, IdentifierType, Message, @@ -418,6 +421,27 @@ def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece ``not_in_memory``), each carrying a non-empty ``conversation_id``. """ + def _persist_message_pieces(self, *, message_pieces: Sequence[MessagePiece]) -> None: + entries = [PromptMemoryEntry(entry=piece) for piece in message_pieces] + with closing(self.get_session()) as session: + try: + for piece, entry in zip(message_pieces, entries, strict=True): + for position, identifier in enumerate(piece.converter_identifiers): + converter_identifier = ConverterIdentifier.from_component_identifier(identifier) + self._persist_identifier(session=session, identifier=converter_identifier) + entry.converter_identifier_links.append( + PromptConverterIdentifierEntry( + position=position, + converter_identifier_hash=converter_identifier.hash, + ) + ) + session.add_all(entries) + session.commit() + except SQLAlchemyError as e: + session.rollback() + logger.exception(f"Error inserting prompt memory entries: {e}") + raise + @staticmethod def _validate_persistable_conversation_ids(*, message_pieces: Sequence[MessagePiece]) -> None: """ @@ -535,6 +559,8 @@ def _persist_identifier(cls, *, session: Any, identifier: ComponentIdentifier) - def _get_identifier_entry_type(identifier: ComponentIdentifier) -> type[ComponentIdentifierEntry[Any]]: if isinstance(identifier, TargetIdentifier): return TargetIdentifierEntry + if isinstance(identifier, ConverterIdentifier): + return ConverterIdentifierEntry if isinstance(identifier, ScorerIdentifier): return ScorerIdentifierEntry if isinstance(identifier, ScenarioIdentifier): diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 60f927b9d1..9bbb7bd436 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -45,6 +45,7 @@ Conversation, ConversationReference, ConversationType, + ConverterIdentifier, EvaluationIdentifier, MessagePiece, PromptDataType, @@ -283,6 +284,13 @@ class PromptMemoryEntry(Base): back_populates="prompt_request_piece", foreign_keys="ScoreEntry.prompt_request_response_id", ) + converter_identifier_links: Mapped[list["PromptConverterIdentifierEntry"]] = relationship( + "PromptConverterIdentifierEntry", + primaryjoin="PromptMemoryEntry.id == PromptConverterIdentifierEntry.prompt_memory_entry_id", + foreign_keys="PromptConverterIdentifierEntry.prompt_memory_entry_id", + order_by="PromptConverterIdentifierEntry.position", + cascade="all, delete-orphan", + ) def __init__(self, *, entry: MessagePiece) -> None: """ @@ -622,6 +630,60 @@ class TargetIdentifierChildEntry(Base): ) +class ConverterIdentifierEntry(ComponentIdentifierEntry[ConverterIdentifier]): + """Content-addressed store of ``ConverterIdentifier`` projections.""" + + __tablename__ = "ConverterIdentifiers" + __table_args__ = {"extend_existing": True} + + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = { + "converter_target": "converter_target_hash", + "sub_converter": "sub_converter_hash", + } + + supported_input_types: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + supported_output_types: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + converter_target_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) + sub_converter_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey("ConverterIdentifiers.hash"), nullable=True + ) + + converter_target: Mapped["TargetIdentifierEntry | None"] = relationship( + "TargetIdentifierEntry", + foreign_keys=[converter_target_hash], + ) + sub_converter: Mapped["ConverterIdentifierEntry | None"] = relationship( + "ConverterIdentifierEntry", + foreign_keys=[sub_converter_hash], + remote_side="ConverterIdentifierEntry.hash", + ) + + +class PromptConverterIdentifierEntry(Base): + """Ordered association between a prompt piece and an applied converter.""" + + __tablename__ = "PromptConverterIdentifiers" + __table_args__ = {"extend_existing": True} + + prompt_memory_entry_id: Mapped[uuid.UUID] = mapped_column( + CustomUUID, + ForeignKey(f"{PromptMemoryEntry.__tablename__}.id"), + primary_key=True, + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + converter_identifier_hash: Mapped[str] = mapped_column( + String(64), + ForeignKey(f"{ConverterIdentifierEntry.__tablename__}.hash"), + nullable=False, + ) + converter_identifier: Mapped["ConverterIdentifierEntry"] = relationship( + "ConverterIdentifierEntry", + foreign_keys=[converter_identifier_hash], + ) + + class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): """Content-addressed store of ``ScorerIdentifier`` projections.""" diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index a3bbf8c119..a2f4139550 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -28,7 +28,7 @@ ScenarioResultEntry, ) from pyrit.memory.storage import DiskStorageIO -from pyrit.models import ConversationStats, MessagePiece +from pyrit.models import ConversationStats logger = logging.getLogger(__name__) @@ -292,16 +292,6 @@ def _get_condition_json_array_match( combined = joiner.join(conditions) return text(f"({combined})").bindparams(**bindparams_dict) - def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece]) -> None: - """ - Persist already-validated message pieces to the SQLite store. - - Args: - message_pieces (Sequence[MessagePiece]): Persistable pieces (filtered and - validated by ``add_message_pieces_to_memory``). - """ - self._insert_entries(entries=[PromptMemoryEntry(entry=piece) for piece in message_pieces]) - def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ Insert embedding data into memory storage. diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index 85d934e923..d1fec98ebf 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -13,16 +13,23 @@ from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory import MemoryInterface, PromptMemoryEntry +from pyrit.memory.memory_models import ( + ConverterIdentifierEntry, + PromptConverterIdentifierEntry, + TargetIdentifierEntry, +) from pyrit.memory.storage.serializers import set_message_piece_sha256_async from pyrit.models import ( ComponentIdentifier, Conversation, + ConverterIdentifier, IdentifierFilter, IdentifierType, Message, MessagePiece, Score, SeedPrompt, + TargetIdentifier, ) @@ -62,6 +69,51 @@ def test_add_message_pieces_to_memory( assert len(sqlite_instance.get_message_pieces()) == num_conversations +def test_add_message_pieces_persists_converter_identifier_graph(sqlite_instance: MemoryInterface): + target = TargetIdentifier( + class_name="ConverterTarget", + class_module="tests.unit.memory", + model_name="converter-model", + ) + nested = ConverterIdentifier( + class_name="NestedConverter", + class_module="tests.unit.memory", + supported_input_types=["text"], + supported_output_types=["text"], + converter_target=target, + ) + converter = ConverterIdentifier( + class_name="CompositeConverter", + class_module="tests.unit.memory", + supported_input_types=["text"], + supported_output_types=["text"], + sub_converter=nested, + ) + pieces = [ + MessagePiece( + conversation_id=str(uuid4()), + role="user", + original_value=f"Prompt {index}", + converter_identifiers=[converter, nested], + ) + for index in range(2) + ] + + sqlite_instance.add_message_pieces_to_memory(message_pieces=pieces) + + converter_rows = sqlite_instance._query_entries(ConverterIdentifierEntry) + link_rows = sqlite_instance._query_entries(PromptConverterIdentifierEntry) + assert {row.hash for row in converter_rows} == {converter.hash, nested.hash} + assert next(row for row in converter_rows if row.hash == converter.hash).sub_converter_hash == nested.hash + assert next(row for row in converter_rows if row.hash == nested.hash).converter_target_hash == target.hash + assert len(sqlite_instance._query_entries(TargetIdentifierEntry)) == 1 + assert len(link_rows) == 4 + assert {(row.position, row.converter_identifier_hash) for row in link_rows} == { + (0, converter.hash), + (1, nested.hash), + } + + def test_get_message_pieces_uuid_and_string_ids(sqlite_instance: MemoryInterface): """Test that get_message_pieces handles both UUID objects and string representations.""" uuid1 = uuid.uuid4() diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 05bc00a65e..774141d01f 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -14,6 +14,7 @@ AttackResultEntry, ComponentIdentifierEntry, ConversationMessageWithSimilarity, + ConverterIdentifierEntry, EmbeddingDataEntry, EmbeddingMessageWithSimilarity, PromptMemoryEntry, @@ -33,6 +34,7 @@ ComponentIdentifier, ConversationReference, ConversationType, + ConverterIdentifier, MessagePiece, ScenarioIdentifier, ScenarioResult, @@ -204,6 +206,7 @@ def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): ("identifier_type", "entry_type"), [ (TargetIdentifier, TargetIdentifierEntry), + (ConverterIdentifier, ConverterIdentifierEntry), (ScorerIdentifier, ScorerIdentifierEntry), (ScenarioIdentifier, ScenarioIdentifierEntry), ], diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index ef1c04e59b..092d7b75c7 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -877,6 +877,80 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): engine.dispose() +# ============================================================================= +# Backfill tests for converter identifier persistence (c8e1f3a5b7d9) +# ============================================================================= + + +def test_converter_identifier_migration_backfills_graph_and_prompt_links(): + """Existing converter JSON is normalized with dependencies and ordered prompt links.""" + from pyrit.models import ConverterIdentifier, TargetIdentifier + + target = TargetIdentifier( + class_name="ConverterTarget", + class_module="pyrit.prompt_target", + model_name="converter-model", + ) + nested = ConverterIdentifier( + class_name="NestedConverter", + class_module="pyrit.prompt_converter", + supported_input_types=["text"], + supported_output_types=["text"], + converter_target=target, + ) + converter = ConverterIdentifier( + class_name="CompositeConverter", + class_module="pyrit.prompt_converter", + supported_input_types=["text"], + supported_output_types=["text"], + sub_converter=nested, + ) + prompt_id = uuid.uuid4() + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'converter-backfill.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "b7d9f1a3c5e7") + connection.execute( + text( + 'INSERT INTO "PromptMemoryEntries" ' + "(id, role, conversation_id, sequence, timestamp, labels, prompt_metadata, " + "converter_identifiers, original_value_data_type, original_value, " + "converted_value_data_type, original_prompt_id, pyrit_version) " + "VALUES (:id, 'user', 'conversation', 0, '2026-07-13', '{}', '{}', :identifiers, " + "'text', 'prompt', 'text', :id, '0.10.0')" + ), + { + "id": str(prompt_id), + "identifiers": json.dumps([converter.model_dump(), nested.model_dump()]), + }, + ) + + command.upgrade(config, "c8e1f3a5b7d9") + + converter_rows = connection.execute( + text('SELECT hash, converter_target_hash, sub_converter_hash FROM "ConverterIdentifiers"') + ).fetchall() + target_hashes = connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars().all() + links = connection.execute( + text( + 'SELECT position, converter_identifier_hash FROM "PromptConverterIdentifiers" ORDER BY position' + ) + ).fetchall() + + assert {row[0] for row in converter_rows} == {converter.hash, nested.hash} + root_row = next(row for row in converter_rows if row[0] == converter.hash) + nested_row = next(row for row in converter_rows if row[0] == nested.hash) + assert root_row[2] == nested.hash + assert nested_row[1] == target.hash + assert target_hashes == [target.hash] + assert links == [(0, converter.hash), (1, nested.hash)] + finally: + engine.dispose() + + # ============================================================================= # Backfill tests for scenario identifier persistence (b7d9f1a3c5e7) # ============================================================================= From 49be08211945458310f0a3f794e26e32cdf6d5a6 Mon Sep 17 00:00:00 2001 From: Behnam Ousat Date: Mon, 13 Jul 2026 17:03:08 -0700 Subject: [PATCH 08/24] persist attack identifiers --- ...2a4b6c8e0_add_attack_identifiers_tables.py | 401 ++++++++++++++++++ pyrit/memory/memory_interface.py | 24 ++ pyrit/memory/memory_models.py | 232 +++++++++- tests/unit/memory/test_memory_models.py | 94 ++++ tests/unit/memory/test_migration.py | 115 +++++ 5 files changed, 859 insertions(+), 7 deletions(-) create mode 100644 pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py diff --git a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py new file mode 100644 index 0000000000..426e8295d7 --- /dev/null +++ b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py @@ -0,0 +1,401 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist attack identifier graphs as content-addressed rows. + +Revision ID: d9f2a4b6c8e0 +Revises: c8e1f3a5b7d9 +Create Date: 2026-07-13 17:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 +from typing import Any + +import sqlalchemy as sa +from alembic import op + +revision: str = "d9f2a4b6c8e0" +down_revision: str | None = "c8e1f3a5b7d9" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + _create_identifier_tables() + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.add_column(sa.Column("atomic_attack_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_attack_result_entries_atomic_attack_identifier_hash", + "AtomicAttackIdentifiers", + ["atomic_attack_identifier_hash"], + ["hash"], + ) + _backfill_attack_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.drop_column("atomic_attack_identifier_hash") + op.drop_table("AtomicAttackSeedIdentifiers") + op.drop_table("AtomicAttackIdentifiers") + op.drop_table("AttackTechniqueSeedIdentifiers") + op.drop_table("AttackTechniqueIdentifiers") + op.drop_table("AttackResponseConverterIdentifiers") + op.drop_table("AttackRequestConverterIdentifiers") + op.drop_table("AttackIdentifiers") + op.drop_table("SeedIdentifiers") + + +def _create_identifier_tables() -> None: + op.create_table( + "SeedIdentifiers", + *_common_columns(), + sa.Column("value", sa.Unicode(), nullable=True), + sa.Column("value_sha256", sa.String(), nullable=True), + sa.Column("data_type", sa.String(), nullable=True), + sa.Column("dataset_name", sa.String(), nullable=True), + sa.Column("is_general_technique", sa.Boolean(), nullable=True), + ) + op.create_table( + "AttackIdentifiers", + *_common_columns(), + sa.Column("adversarial_system_prompt", sa.Unicode(), nullable=True), + sa.Column("adversarial_seed_prompt", sa.Unicode(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("adversarial_chat_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["objective_target_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["adversarial_chat_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["objective_scorer_hash"], ["ScorerIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + _create_ordered_edge_table( + table_name="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + op.create_table( + "AttackTechniqueIdentifiers", + *_common_columns(), + sa.Column("attack_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["attack_identifier_hash"], ["AttackIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_table="AttackTechniqueIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + op.create_table( + "AtomicAttackIdentifiers", + *_common_columns(), + sa.Column("attack_technique_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["attack_technique_identifier_hash"], + ["AttackTechniqueIdentifiers.hash"], + ), + ) + _create_ordered_edge_table( + table_name="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_table="AtomicAttackIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + + +def _common_columns() -> tuple[sa.Column[Any], ...]: + return ( + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + +def _create_ordered_edge_table( + *, + table_name: str, + parent_column: str, + parent_table: str, + child_column: str, + child_table: str, +) -> None: + op.create_table( + table_name, + sa.Column(parent_column, sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column(child_column, sa.String(64), nullable=False), + sa.ForeignKeyConstraint([parent_column], [f"{parent_table}.hash"]), + sa.ForeignKeyConstraint([child_column], [f"{child_table}.hash"]), + sa.PrimaryKeyConstraint(parent_column, "position"), + ) + + +def _backfill_attack_identifiers() -> None: + """Backfill attack identifier graphs and result links from retained JSON.""" + from pyrit.models import AtomicAttackIdentifier + + bind = op.get_bind() + result_rows = bind.execute( + sa.text( + 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" ' + "WHERE atomic_attack_identifier IS NOT NULL" + ) + ).fetchall() + inserter = _AttackIdentifierGraphInserter(bind=bind) + update_stmt = sa.text( + 'UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id' + ) + + skipped = 0 + for result_id, raw_identifier in result_rows: + try: + stored = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier + identifier = AtomicAttackIdentifier.model_validate(stored) + inserter.insert_atomic_attack(identifier) + bind.execute(update_stmt, {"hash": identifier.hash, "id": result_id}) + except Exception: + skipped += 1 + logger.warning( + f"Attack identifier backfill could not reconstruct result {result_id}", + exc_info=True, + ) + if skipped: + logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") + + +class _AttackIdentifierGraphInserter: + """Insert a normalized attack identifier graph during migration backfill.""" + + def __init__(self, *, bind: Any) -> None: + self._bind = bind + self._hashes = { + table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) + for table in ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ConverterIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ) + } + + def insert_atomic_attack(self, identifier: Any) -> None: + if identifier.hash in self._hashes["AtomicAttackIdentifiers"]: + return + if identifier.attack_technique is not None: + self._insert_attack_technique(identifier.attack_technique) + for seed in identifier.seed_identifiers: + self._insert_seed(seed) + self._insert_identifier( + table="AtomicAttackIdentifiers", + identifier=identifier, + extra={ + "attack_technique_identifier_hash": ( + identifier.attack_technique.hash if identifier.attack_technique is not None else None + ) + }, + ) + self._insert_edges( + table="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_hash=identifier.hash, + child_column="seed_identifier_hash", + children=identifier.seed_identifiers, + ) + + def _insert_attack_technique(self, identifier: Any) -> None: + if identifier.hash in self._hashes["AttackTechniqueIdentifiers"]: + return + if identifier.attack is not None: + self._insert_attack(identifier.attack) + for seed in identifier.technique_seeds: + self._insert_seed(seed) + self._insert_identifier( + table="AttackTechniqueIdentifiers", + identifier=identifier, + extra={"attack_identifier_hash": identifier.attack.hash if identifier.attack is not None else None}, + ) + self._insert_edges( + table="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_hash=identifier.hash, + child_column="seed_identifier_hash", + children=identifier.technique_seeds, + ) + + def _insert_attack(self, identifier: Any) -> None: + if identifier.hash in self._hashes["AttackIdentifiers"]: + return + for target in (identifier.objective_target, identifier.adversarial_chat): + if target is not None: + self._insert_target(target) + if identifier.objective_scorer is not None: + self._insert_scorer(identifier.objective_scorer) + for converter in [*identifier.request_converters, *identifier.response_converters]: + self._insert_converter(converter) + self._insert_identifier( + table="AttackIdentifiers", + identifier=identifier, + extra={ + "adversarial_system_prompt": identifier.adversarial_system_prompt, + "adversarial_seed_prompt": identifier.adversarial_seed_prompt, + "objective_target_hash": ( + identifier.objective_target.hash if identifier.objective_target is not None else None + ), + "adversarial_chat_hash": ( + identifier.adversarial_chat.hash if identifier.adversarial_chat is not None else None + ), + "objective_scorer_hash": ( + identifier.objective_scorer.hash if identifier.objective_scorer is not None else None + ), + }, + ) + self._insert_edges( + table="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier.hash, + child_column="converter_identifier_hash", + children=identifier.request_converters, + ) + self._insert_edges( + table="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier.hash, + child_column="converter_identifier_hash", + children=identifier.response_converters, + ) + + def _insert_seed(self, identifier: Any) -> None: + if identifier.hash in self._hashes["SeedIdentifiers"]: + return + self._insert_identifier( + table="SeedIdentifiers", + identifier=identifier, + extra=identifier.promoted_scalar_values(), + ) + + def _insert_target(self, identifier: Any) -> None: + if identifier.hash in self._hashes["TargetIdentifiers"]: + return + for child in identifier.targets: + self._insert_target(child) + self._insert_identifier( + table="TargetIdentifiers", + identifier=identifier, + extra=identifier.promoted_scalar_values(), + ) + self._insert_edges( + table="TargetIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier.hash, + child_column="child_hash", + children=identifier.targets, + ) + + def _insert_scorer(self, identifier: Any) -> None: + if identifier.hash in self._hashes["ScorerIdentifiers"]: + return + if identifier.prompt_target is not None: + self._insert_target(identifier.prompt_target) + for child in identifier.sub_scorers: + self._insert_scorer(child) + self._insert_identifier( + table="ScorerIdentifiers", + identifier=identifier, + extra={ + **identifier.promoted_scalar_values(), + "prompt_target_hash": ( + identifier.prompt_target.hash if identifier.prompt_target is not None else None + ), + }, + ) + self._insert_edges( + table="ScorerIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier.hash, + child_column="child_hash", + children=identifier.sub_scorers, + ) + + def _insert_converter(self, identifier: Any) -> None: + if identifier.hash in self._hashes["ConverterIdentifiers"]: + return + if identifier.converter_target is not None: + self._insert_target(identifier.converter_target) + if identifier.sub_converter is not None: + self._insert_converter(identifier.sub_converter) + self._insert_identifier( + table="ConverterIdentifiers", + identifier=identifier, + extra={ + **identifier.promoted_scalar_values(), + "converter_target_hash": ( + identifier.converter_target.hash if identifier.converter_target is not None else None + ), + "sub_converter_hash": ( + identifier.sub_converter.hash if identifier.sub_converter is not None else None + ), + }, + ) + + def _insert_identifier(self, *, table: str, identifier: Any, extra: dict[str, Any]) -> None: + columns = ["hash", "class_name", "class_module", "identifier_json", *extra, "pyrit_version"] + placeholders = [f":{column}" for column in columns] + values = { + "hash": identifier.hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "pyrit_version": identifier.pyrit_version, + **{ + name: json.dumps(value) if isinstance(value, (list, dict)) else value + for name, value in extra.items() + }, + } + self._bind.execute( + sa.text(f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(placeholders)})'), + values, + ) + self._hashes[table].add(identifier.hash) + + def _insert_edges( + self, + *, + table: str, + parent_column: str, + parent_hash: str, + child_column: str, + children: Sequence[Any], + ) -> None: + statement = sa.text( + f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' + f'VALUES (:parent_hash, :position, :child_hash)' + ) + for position, child in enumerate(children): + self._bind.execute( + statement, + {"parent_hash": parent_hash, "position": position, "child_hash": child.hash}, + ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 8c6eb0b152..8a70e65d98 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -21,7 +21,10 @@ from pyrit.memory.memory_embedding import MemoryEmbedding from pyrit.memory.memory_models import ( + AtomicAttackIdentifierEntry, + AttackIdentifierEntry, AttackResultEntry, + AttackTechniqueIdentifierEntry, Base, ComponentIdentifierEntry, ConversationEntry, @@ -34,6 +37,7 @@ ScoreEntry, ScorerIdentifierEntry, SeedEntry, + SeedIdentifierEntry, TargetIdentifierEntry, ) from pyrit.memory.storage import ( @@ -43,7 +47,10 @@ set_seed_sha256_async, ) from pyrit.models import ( + AtomicAttackIdentifier, + AttackIdentifier, AttackResult, + AttackTechniqueIdentifier, ComponentIdentifier, Conversation, ConversationStats, @@ -59,6 +66,7 @@ Seed, SeedDataset, SeedGroup, + SeedIdentifier, SeedType, TargetIdentifier, group_conversation_message_pieces_by_sequence, @@ -557,6 +565,14 @@ def _persist_identifier(cls, *, session: Any, identifier: ComponentIdentifier) - @staticmethod def _get_identifier_entry_type(identifier: ComponentIdentifier) -> type[ComponentIdentifierEntry[Any]]: + if isinstance(identifier, AtomicAttackIdentifier): + return AtomicAttackIdentifierEntry + if isinstance(identifier, AttackTechniqueIdentifier): + return AttackTechniqueIdentifierEntry + if isinstance(identifier, AttackIdentifier): + return AttackIdentifierEntry + if isinstance(identifier, SeedIdentifier): + return SeedIdentifierEntry if isinstance(identifier, TargetIdentifier): return TargetIdentifierEntry if isinstance(identifier, ConverterIdentifier): @@ -1893,6 +1909,14 @@ def add_attack_results_to_memory(self, *, attack_results: Sequence[AttackResult] entries = [AttackResultEntry(entry=attack_result) for attack_result in attack_results] with closing(self.get_session()) as session: try: + for attack_result in attack_results: + if attack_result.atomic_attack_identifier is not None: + self._persist_identifier( + session=session, + identifier=AtomicAttackIdentifier.from_component_identifier( + attack_result.atomic_attack_identifier + ), + ) session.add_all(entries) session.commit() except SQLAlchemyError: diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 9bbb7bd436..cdd7948985 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -15,6 +15,7 @@ ARRAY, INTEGER, JSON, + Boolean, DateTime, Float, ForeignKey, @@ -38,8 +39,11 @@ from pyrit.models import ( SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY, AtomicAttackEvaluationIdentifier, + AtomicAttackIdentifier, + AttackIdentifier, AttackOutcome, AttackResult, + AttackTechniqueIdentifier, ChatMessageRole, ComponentIdentifier, Conversation, @@ -57,6 +61,7 @@ ScorerEvaluationIdentifier, ScorerIdentifier, Seed, + SeedIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, @@ -430,7 +435,8 @@ class _ChildRelationshipSpec: relationship_name: str edge_factory: Callable[[], Any] - edge_child_attr: str + edge_child_attr: str | None = None + edge_child_hash_attr: str | None = None edge_position_attr: str | None = "position" @@ -530,10 +536,13 @@ def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: for position, child_identifier in enumerate(children): if not isinstance(child_identifier, ComponentIdentifier): continue - typed_child_identifier = cast("T", child_identifier) - child_entry = cls._from_domain_model_recursive(domain_model=typed_child_identifier, seen=seen) edge_row = spec.edge_factory() - setattr(edge_row, spec.edge_child_attr, child_entry) + if spec.edge_child_hash_attr is not None: + setattr(edge_row, spec.edge_child_hash_attr, child_identifier.hash) + elif spec.edge_child_attr is not None: + typed_child_identifier = cast("T", child_identifier) + child_entry = cls._from_domain_model_recursive(domain_model=typed_child_identifier, seen=seen) + setattr(edge_row, spec.edge_child_attr, child_entry) if spec.edge_position_attr: setattr(edge_row, spec.edge_position_attr, position) edge_rows.append(edge_row) @@ -775,6 +784,204 @@ class ScenarioIdentifierEntry(ComponentIdentifierEntry[ScenarioIdentifier]): ) +class SeedIdentifierEntry(ComponentIdentifierEntry[SeedIdentifier]): + """Content-addressed store of ``SeedIdentifier`` projections.""" + + __tablename__ = "SeedIdentifiers" + __table_args__ = {"extend_existing": True} + + value: Mapped[str | None] = mapped_column(Unicode, nullable=True) + value_sha256: Mapped[str | None] = mapped_column(String, nullable=True) + data_type: Mapped[str | None] = mapped_column(String, nullable=True) + dataset_name: Mapped[str | None] = mapped_column(String, nullable=True) + is_general_technique: Mapped[bool | None] = mapped_column(Boolean, nullable=True) + + +class AttackIdentifierEntry(ComponentIdentifierEntry[AttackIdentifier]): + """Content-addressed store of ``AttackIdentifier`` projections.""" + + __tablename__ = "AttackIdentifiers" + __table_args__ = {"extend_existing": True} + + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = { + "request_converters": _ChildRelationshipSpec( + relationship_name="request_converters", + edge_factory=lambda: AttackRequestConverterIdentifierEntry(), + edge_child_hash_attr="converter_identifier_hash", + ), + "response_converters": _ChildRelationshipSpec( + relationship_name="response_converters", + edge_factory=lambda: AttackResponseConverterIdentifierEntry(), + edge_child_hash_attr="converter_identifier_hash", + ), + } + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = { + "objective_target": "objective_target_hash", + "adversarial_chat": "adversarial_chat_hash", + "objective_scorer": "objective_scorer_hash", + } + + adversarial_system_prompt: Mapped[str | None] = mapped_column(Unicode, nullable=True) + adversarial_seed_prompt: Mapped[str | None] = mapped_column(Unicode, nullable=True) + objective_target_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) + adversarial_chat_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) + objective_scorer_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{ScorerIdentifierEntry.__tablename__}.hash"), nullable=True + ) + + objective_target: Mapped["TargetIdentifierEntry | None"] = relationship( + "TargetIdentifierEntry", foreign_keys=[objective_target_hash] + ) + adversarial_chat: Mapped["TargetIdentifierEntry | None"] = relationship( + "TargetIdentifierEntry", foreign_keys=[adversarial_chat_hash] + ) + objective_scorer: Mapped["ScorerIdentifierEntry | None"] = relationship( + "ScorerIdentifierEntry", foreign_keys=[objective_scorer_hash] + ) + request_converters: Mapped[list["AttackRequestConverterIdentifierEntry"]] = relationship( + "AttackRequestConverterIdentifierEntry", + order_by="AttackRequestConverterIdentifierEntry.position", + cascade="all, delete-orphan", + ) + response_converters: Mapped[list["AttackResponseConverterIdentifierEntry"]] = relationship( + "AttackResponseConverterIdentifierEntry", + order_by="AttackResponseConverterIdentifierEntry.position", + cascade="all, delete-orphan", + ) + + +class AttackRequestConverterIdentifierEntry(Base): + """Ordered request-converter edge for an attack identifier.""" + + __tablename__ = "AttackRequestConverterIdentifiers" + __table_args__ = {"extend_existing": True} + + attack_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{AttackIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + converter_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{ConverterIdentifierEntry.__tablename__}.hash"), nullable=False + ) + converter_identifier: Mapped["ConverterIdentifierEntry"] = relationship( + "ConverterIdentifierEntry", foreign_keys=[converter_identifier_hash] + ) + + +class AttackResponseConverterIdentifierEntry(Base): + """Ordered response-converter edge for an attack identifier.""" + + __tablename__ = "AttackResponseConverterIdentifiers" + __table_args__ = {"extend_existing": True} + + attack_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{AttackIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + converter_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{ConverterIdentifierEntry.__tablename__}.hash"), nullable=False + ) + converter_identifier: Mapped["ConverterIdentifierEntry"] = relationship( + "ConverterIdentifierEntry", foreign_keys=[converter_identifier_hash] + ) + + +class AttackTechniqueIdentifierEntry(ComponentIdentifierEntry[AttackTechniqueIdentifier]): + """Content-addressed store of ``AttackTechniqueIdentifier`` projections.""" + + __tablename__ = "AttackTechniqueIdentifiers" + __table_args__ = {"extend_existing": True} + + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = { + "technique_seeds": _ChildRelationshipSpec( + relationship_name="technique_seeds", + edge_factory=lambda: AttackTechniqueSeedIdentifierEntry(), + edge_child_hash_attr="seed_identifier_hash", + ) + } + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = {"attack": "attack_identifier_hash"} + + attack_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{AttackIdentifierEntry.__tablename__}.hash"), nullable=True + ) + attack: Mapped["AttackIdentifierEntry | None"] = relationship( + "AttackIdentifierEntry", foreign_keys=[attack_identifier_hash] + ) + technique_seeds: Mapped[list["AttackTechniqueSeedIdentifierEntry"]] = relationship( + "AttackTechniqueSeedIdentifierEntry", + order_by="AttackTechniqueSeedIdentifierEntry.position", + cascade="all, delete-orphan", + ) + + +class AttackTechniqueSeedIdentifierEntry(Base): + """Ordered seed edge for an attack technique identifier.""" + + __tablename__ = "AttackTechniqueSeedIdentifiers" + __table_args__ = {"extend_existing": True} + + attack_technique_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{AttackTechniqueIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + seed_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{SeedIdentifierEntry.__tablename__}.hash"), nullable=False + ) + seed_identifier: Mapped["SeedIdentifierEntry"] = relationship( + "SeedIdentifierEntry", foreign_keys=[seed_identifier_hash] + ) + + +class AtomicAttackIdentifierEntry(ComponentIdentifierEntry[AtomicAttackIdentifier]): + """Content-addressed store of ``AtomicAttackIdentifier`` projections.""" + + __tablename__ = "AtomicAttackIdentifiers" + __table_args__ = {"extend_existing": True} + + CHILD_RELATIONSHIP_SPECS: ClassVar[dict[str, _ChildRelationshipSpec]] = { + "seed_identifiers": _ChildRelationshipSpec( + relationship_name="seed_identifiers", + edge_factory=lambda: AtomicAttackSeedIdentifierEntry(), + edge_child_hash_attr="seed_identifier_hash", + ) + } + CHILD_HASH_COLUMNS: ClassVar[dict[str, str]] = {"attack_technique": "attack_technique_identifier_hash"} + + attack_technique_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{AttackTechniqueIdentifierEntry.__tablename__}.hash"), nullable=True + ) + attack_technique: Mapped["AttackTechniqueIdentifierEntry | None"] = relationship( + "AttackTechniqueIdentifierEntry", foreign_keys=[attack_technique_identifier_hash] + ) + seed_identifiers: Mapped[list["AtomicAttackSeedIdentifierEntry"]] = relationship( + "AtomicAttackSeedIdentifierEntry", + order_by="AtomicAttackSeedIdentifierEntry.position", + cascade="all, delete-orphan", + ) + + +class AtomicAttackSeedIdentifierEntry(Base): + """Ordered seed edge for an atomic attack identifier.""" + + __tablename__ = "AtomicAttackSeedIdentifiers" + __table_args__ = {"extend_existing": True} + + atomic_attack_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{AtomicAttackIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + seed_identifier_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{SeedIdentifierEntry.__tablename__}.hash"), nullable=False + ) + seed_identifier: Mapped["SeedIdentifierEntry"] = relationship( + "SeedIdentifierEntry", foreign_keys=[seed_identifier_hash] + ) + + class ConversationEntry(Base): """ Conversation-scoped metadata, persisted once per ``conversation_id``. @@ -1294,6 +1501,9 @@ class AttackResultEntry(Base): conversation_id = mapped_column(String, nullable=False) objective = mapped_column(Unicode, nullable=False) atomic_attack_identifier: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) + atomic_attack_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{AtomicAttackIdentifierEntry.__tablename__}.hash"), nullable=True + ) objective_sha256 = mapped_column(String, nullable=True) last_response_id: Mapped[uuid.UUID | None] = mapped_column( CustomUUID, ForeignKey(f"{PromptMemoryEntry.__tablename__}.id"), nullable=True @@ -1347,6 +1557,10 @@ class AttackResultEntry(Base): "ScoreEntry", foreign_keys=[last_score_id], ) + atomic_attack_identifier_entry: Mapped["AtomicAttackIdentifierEntry | None"] = relationship( + "AtomicAttackIdentifierEntry", + foreign_keys=[atomic_attack_identifier_hash], + ) def __init__(self, *, entry: AttackResult) -> None: """ @@ -1360,13 +1574,17 @@ def __init__(self, *, entry: AttackResult) -> None: self.objective = entry.objective # Always recompute eval_hash before dumping so the stored JSON carries the # freshly computed value for DB-level filtering (never a value from storage). + atomic_attack_identifier = None if entry.atomic_attack_identifier: - entry.atomic_attack_identifier = entry.atomic_attack_identifier.with_eval_hash( - AtomicAttackEvaluationIdentifier(entry.atomic_attack_identifier).eval_hash + atomic_attack_identifier = AtomicAttackIdentifier.from_component_identifier(entry.atomic_attack_identifier) + atomic_attack_identifier = atomic_attack_identifier.with_eval_hash( + AtomicAttackEvaluationIdentifier(atomic_attack_identifier).eval_hash ) + entry.atomic_attack_identifier = atomic_attack_identifier self.atomic_attack_identifier = ( - entry.atomic_attack_identifier.model_dump() if entry.atomic_attack_identifier else None + atomic_attack_identifier.model_dump() if atomic_attack_identifier else None ) + self.atomic_attack_identifier_hash = atomic_attack_identifier.hash if atomic_attack_identifier else None self.objective_sha256 = to_sha256(entry.objective) # Use helper method for UUID conversions diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 774141d01f..759d5d9424 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -9,9 +9,19 @@ import pytest from pydantic import ValidationError +from sqlalchemy import create_engine, select +from sqlalchemy.orm import Session from pyrit.memory.memory_models import ( + AtomicAttackIdentifierEntry, + AtomicAttackSeedIdentifierEntry, + AttackIdentifierEntry, + AttackRequestConverterIdentifierEntry, + AttackResponseConverterIdentifierEntry, AttackResultEntry, + AttackTechniqueIdentifierEntry, + AttackTechniqueSeedIdentifierEntry, + Base, ComponentIdentifierEntry, ConversationMessageWithSimilarity, ConverterIdentifierEntry, @@ -23,14 +33,17 @@ ScoreEntry, ScorerIdentifierEntry, SeedEntry, + SeedIdentifierEntry, TargetIdentifierEntry, UTCDateTime, _load_identifier, ) from pyrit.models import ( AtomicAttackIdentifier, + AttackIdentifier, AttackOutcome, AttackResult, + AttackTechniqueIdentifier, ComponentIdentifier, ConversationReference, ConversationType, @@ -40,6 +53,7 @@ ScenarioResult, Score, ScorerIdentifier, + SeedIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, @@ -209,6 +223,10 @@ def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): (ConverterIdentifier, ConverterIdentifierEntry), (ScorerIdentifier, ScorerIdentifierEntry), (ScenarioIdentifier, ScenarioIdentifierEntry), + (SeedIdentifier, SeedIdentifierEntry), + (AttackIdentifier, AttackIdentifierEntry), + (AttackTechniqueIdentifier, AttackTechniqueIdentifierEntry), + (AtomicAttackIdentifier, AtomicAttackIdentifierEntry), ], ) def test_identifier_entry_maps_promoted_children_by_cardinality( @@ -227,6 +245,82 @@ def test_identifier_entry_maps_promoted_children_by_cardinality( assert set(entry_type.CHILD_HASH_COLUMNS) == singular_children +def test_atomic_attack_identifier_graph_persists_with_result_link() -> None: + target = TargetIdentifier(class_name="Target", class_module="pyrit.prompt_target", model_name="model") + scorer = ScorerIdentifier(class_name="Scorer", class_module="pyrit.score", scorer_type="true_false") + converter = ConverterIdentifier( + class_name="Converter", + class_module="pyrit.prompt_converter", + supported_input_types=["text"], + supported_output_types=["text"], + ) + technique_seed = SeedIdentifier( + class_name="Seed", + class_module="pyrit.models", + value="technique seed", + data_type="text", + ) + dataset_seed = SeedIdentifier( + class_name="Seed", + class_module="pyrit.models", + value="dataset seed", + data_type="text", + ) + attack = AttackIdentifier( + class_name="Attack", + class_module="pyrit.executor.attack", + objective_target=target, + objective_scorer=scorer, + request_converters=[converter], + response_converters=[converter], + ) + technique = AttackTechniqueIdentifier( + class_name="AttackTechnique", + class_module="pyrit.scenario.core.attack_technique", + attack=attack, + technique_seeds=[technique_seed], + ) + atomic = AtomicAttackIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + attack_technique=technique, + seed_identifiers=[technique_seed, dataset_seed], + ) + result = AttackResult(conversation_id="conversation", objective="objective", atomic_attack_identifier=atomic) + + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(engine) + from pyrit.memory import MemoryInterface + + memory = MagicMock(spec=MemoryInterface) + memory.get_session.side_effect = lambda: Session(engine) + memory._persist_identifier.side_effect = ( + lambda *, session, identifier: MemoryInterface._persist_identifier(session=session, identifier=identifier) + ) + MemoryInterface.add_attack_results_to_memory(memory, attack_results=[result]) + + with Session(engine) as session: + assert session.scalar(select(AttackResultEntry.atomic_attack_identifier_hash)) == atomic.hash + assert session.scalar(select(AtomicAttackIdentifierEntry.hash)) == atomic.hash + assert session.scalar(select(AttackTechniqueIdentifierEntry.hash)) == technique.hash + assert session.scalar(select(AttackIdentifierEntry.hash)) == attack.hash + assert len(session.scalars(select(SeedIdentifierEntry)).all()) == 2 + + technique_edge = session.scalar(select(AttackTechniqueSeedIdentifierEntry)) + assert technique_edge is not None + assert (technique_edge.position, technique_edge.seed_identifier_hash) == (0, technique_seed.hash) + atomic_edges = session.scalars( + select(AtomicAttackSeedIdentifierEntry).order_by(AtomicAttackSeedIdentifierEntry.position) + ).all() + assert [edge.seed_identifier_hash for edge in atomic_edges] == [technique_seed.hash, dataset_seed.hash] + request_edge = session.scalar(select(AttackRequestConverterIdentifierEntry)) + assert request_edge is not None + assert (request_edge.position, request_edge.converter_identifier_hash) == (0, converter.hash) + response_edge = session.scalar(select(AttackResponseConverterIdentifierEntry)) + assert response_edge is not None + assert (response_edge.position, response_edge.converter_identifier_hash) == (0, converter.hash) + + # --------------------------------------------------------------------------- # ConversationMessageWithSimilarity # --------------------------------------------------------------------------- diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 092d7b75c7..17988a8f4f 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -1040,6 +1040,121 @@ def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): engine.dispose() +# ============================================================================= +# Backfill tests for attack identifier persistence (d9f2a4b6c8e0) +# ============================================================================= + + +def test_attack_identifier_migration_backfills_graph_and_result_link(): + """Atomic attack JSON is normalized with dependencies and ordered seed links.""" + from pyrit.models import ( + AtomicAttackIdentifier, + AttackIdentifier, + AttackTechniqueIdentifier, + ConverterIdentifier, + ScorerIdentifier, + SeedIdentifier, + TargetIdentifier, + ) + + target = TargetIdentifier(class_name="Target", class_module="pyrit.prompt_target", model_name="model") + scorer = ScorerIdentifier(class_name="Scorer", class_module="pyrit.score", scorer_type="true_false") + converter = ConverterIdentifier( + class_name="Converter", + class_module="pyrit.prompt_converter", + supported_input_types=["text"], + supported_output_types=["text"], + ) + technique_seed = SeedIdentifier( + class_name="Seed", + class_module="pyrit.models", + value="technique seed", + data_type="text", + ) + dataset_seed = SeedIdentifier( + class_name="Seed", + class_module="pyrit.models", + value="dataset seed", + data_type="text", + ) + attack = AttackIdentifier( + class_name="Attack", + class_module="pyrit.executor.attack", + objective_target=target, + objective_scorer=scorer, + request_converters=[converter], + ) + technique = AttackTechniqueIdentifier( + class_name="AttackTechnique", + class_module="pyrit.scenario.core.attack_technique", + attack=attack, + technique_seeds=[technique_seed], + ) + atomic = AtomicAttackIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + attack_technique=technique, + seed_identifiers=[technique_seed, dataset_seed], + ) + result_id = str(uuid.uuid4()) + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'attack-backfill.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "c8e1f3a5b7d9") + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, atomic_attack_identifier, objective_sha256, " + "executed_turns, execution_time_ms, outcome, timestamp, pyrit_version) " + "VALUES (:id, 'conversation', 'objective', :identifier, 'sha', 1, 0, " + "'success', '2026-07-13', '0.10.0')" + ), + {"id": result_id, "identifier": json.dumps(atomic.model_dump())}, + ) + + command.upgrade(config, "d9f2a4b6c8e0") + + result_hash = connection.execute( + text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() + atomic_row = connection.execute( + text('SELECT hash, attack_technique_identifier_hash FROM "AtomicAttackIdentifiers"') + ).one() + technique_row = connection.execute( + text('SELECT hash, attack_identifier_hash FROM "AttackTechniqueIdentifiers"') + ).one() + attack_row = connection.execute( + text( + "SELECT hash, objective_target_hash, objective_scorer_hash " + 'FROM "AttackIdentifiers"' + ) + ).one() + seed_hashes = set(connection.execute(text('SELECT hash FROM "SeedIdentifiers"')).scalars()) + atomic_seed_hashes = connection.execute( + text( + "SELECT seed_identifier_hash FROM " + '"AtomicAttackSeedIdentifiers" ORDER BY position' + ) + ).scalars().all() + request_converter_hash = connection.execute( + text('SELECT converter_identifier_hash FROM "AttackRequestConverterIdentifiers"') + ).scalar_one() + + assert result_hash == atomic.hash + assert atomic_row == (atomic.hash, technique.hash) + assert technique_row == (technique.hash, attack.hash) + assert attack_row == (attack.hash, target.hash, scorer.hash) + assert seed_hashes == {technique_seed.hash, dataset_seed.hash} + assert atomic_seed_hashes == [technique_seed.hash, dataset_seed.hash] + assert request_converter_hash == converter.hash + finally: + engine.dispose() + + _STRING_TYPES_REQUIRING_LENGTH = {"String", "VARCHAR", "NVARCHAR", "Unicode"} From 6c2a599359509de8183c1f2e0c1063f31f1d4e6f Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 09:23:31 -0700 Subject: [PATCH 09/24] fix integrity check --- pyrit/memory/memory_interface.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 8a70e65d98..8637bf98e5 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -536,10 +536,9 @@ def _persist_target_identifier(cls, *, session: Any, target_identifier: TargetId If the row already exists it was fully persisted before (children and edges included, since rows are immutable), so this returns early. Otherwise the row and - its child edges are inserted inside a savepoint: a concurrent writer that inserts - the same hash first surfaces as an ``IntegrityError`` that rolls back only this - insert (not the surrounding write) and is then treated as a no-op -- the row - exists either way. + its child edges are inserted inside a savepoint. If an ``IntegrityError`` occurs, + it is treated as a concurrent duplicate only when a fresh lookup confirms that + the identifier hash now exists; all other integrity failures are re-raised. Args: session (Any): The active SQLAlchemy session (the caller's transaction). @@ -561,7 +560,10 @@ def _persist_identifier(cls, *, session: Any, identifier: ComponentIdentifier) - session.merge(entry_type.from_domain_model(identifier)) session.flush() except IntegrityError: - pass + with session.no_autoflush: + existing_entry = session.get(entry_type, identifier.hash, populate_existing=True) + if existing_entry is None: + raise @staticmethod def _get_identifier_entry_type(identifier: ComponentIdentifier) -> type[ComponentIdentifierEntry[Any]]: From 6da2994dc232f0e4dcf67aec991343f9f1a8f620 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 09:48:00 -0700 Subject: [PATCH 10/24] fix abstract --- pyrit/memory/memory_interface.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 8637bf98e5..24289a0d9b 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -414,22 +414,17 @@ def add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece] self._validate_persistable_conversation_ids(message_pieces=pieces_to_insert) self._add_message_pieces_to_memory(message_pieces=pieces_to_insert) - @abc.abstractmethod def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece]) -> None: """ Persist already-validated message pieces to the backing store. Called by ``add_message_pieces_to_memory`` after ``not_in_memory`` pieces are - filtered out and conversation_ids are validated. Implementations only translate - the pieces into storage rows and insert them; they must not re-filter or - re-validate. + filtered out and conversation_ids are validated. Args: message_pieces (Sequence[MessagePiece]): Persistable pieces (none flagged ``not_in_memory``), each carrying a non-empty ``conversation_id``. """ - - def _persist_message_pieces(self, *, message_pieces: Sequence[MessagePiece]) -> None: entries = [PromptMemoryEntry(entry=piece) for piece in message_pieces] with closing(self.get_session()) as session: try: From ea4db75c0dbd95d314174e09a17d39bac6959d86 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 10:03:36 -0700 Subject: [PATCH 11/24] fix migrations --- pyrit/memory/alembic/identifier_backfill.py | 359 ++++++++++++++++++ ...c8e0f2b4d6_add_scorer_identifiers_table.py | 115 +----- ...f1a3c5e7_add_scenario_identifiers_table.py | 146 +------ ...3a5b7d9_add_converter_identifiers_table.py | 120 ++---- ...2a4b6c8e0_add_attack_identifiers_tables.py | 248 +----------- ...f7a9c1b3d2_add_target_identifiers_table.py | 106 ++---- pyrit/memory/memory_models.py | 16 +- tests/unit/memory/test_migration.py | 118 ++++++ 8 files changed, 596 insertions(+), 632 deletions(-) create mode 100644 pyrit/memory/alembic/identifier_backfill.py diff --git a/pyrit/memory/alembic/identifier_backfill.py b/pyrit/memory/alembic/identifier_backfill.py new file mode 100644 index 0000000000..a883e7746c --- /dev/null +++ b/pyrit/memory/alembic/identifier_backfill.py @@ -0,0 +1,359 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Frozen, model-independent helpers for identifier migration backfills.""" + +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING, Any + +import sqlalchemy as sa + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + +logger = logging.getLogger(__name__) + + +def run_best_effort_backfill(*, bind: Any, name: str, backfill: Callable[[], None]) -> None: + """Run a data backfill in a savepoint without blocking the schema upgrade.""" + try: + with bind.begin_nested(): + backfill() + except Exception: + logger.warning(f"{name} backfill failed; leaving new identifier links nullable", exc_info=True) + + +def load_identifier(raw_identifier: Any) -> dict[str, Any] | None: + """ + Load a retained identifier JSON value without importing domain models. + + Returns: + dict[str, Any] | None: The identifier dictionary when it has a usable hash. + """ + try: + value = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier + except (TypeError, ValueError): + return None + if not isinstance(value, dict): + return None + identifier_hash = value.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64: + return None + return value + + +def load_identifier_list(raw_identifiers: Any) -> list[dict[str, Any]]: + """ + Load the valid identifiers from a retained JSON list. + + Returns: + list[dict[str, Any]]: Identifier dictionaries carrying usable hashes. + """ + try: + values = json.loads(raw_identifiers) if isinstance(raw_identifiers, str) else raw_identifiers + except (TypeError, ValueError): + return [] + if not isinstance(values, list): + return [] + return [identifier for value in values if (identifier := load_identifier(value)) is not None] + + +class IdentifierGraphInserter: + """Best-effort inserter for the frozen flat identifier JSON shape.""" + + _TABLES = ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ConverterIdentifiers", + "ScenarioIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ) + + def __init__(self, *, bind: Any) -> None: + """Initialize the inserter from tables available at this migration revision.""" + self._bind = bind + table_names = set(sa.inspect(bind).get_table_names()) + self._hashes = { + table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) + for table in self._TABLES + if table in table_names + } + + def insert_target(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a target graph. + + Returns: + str | None: The stored hash when successful. + """ + children = self._children(identifier, "targets") + child_hashes = [child_hash for child in children if (child_hash := self.insert_target(child))] + identifier_hash = self._insert_identifier( + table="TargetIdentifiers", + identifier=identifier, + promoted=( + "endpoint", + "model_name", + "underlying_model_name", + "temperature", + "top_p", + "max_requests_per_minute", + "supported_auth_modes", + ), + ) + if identifier_hash: + self._insert_edges( + table="TargetIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_scorer(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scorer graph. + + Returns: + str | None: The stored hash when successful. + """ + prompt_target = self._child(identifier, "prompt_target", aliases=("chat_target",)) + prompt_target_hash = self.insert_target(prompt_target) if prompt_target else None + sub_scorers = self._children(identifier, "sub_scorers", aliases=("scorers",)) + child_hashes = [child_hash for child in sub_scorers if (child_hash := self.insert_scorer(child))] + identifier_hash = self._insert_identifier( + table="ScorerIdentifiers", + identifier=identifier, + promoted=("scorer_type", "score_aggregator"), + extra={"prompt_target_hash": prompt_target_hash}, + ) + if identifier_hash: + self._insert_edges( + table="ScorerIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_converter(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a converter graph. + + Returns: + str | None: The stored hash when successful. + """ + converter_target = self._child(identifier, "converter_target") + sub_converter = self._child(identifier, "sub_converter") + return self._insert_identifier( + table="ConverterIdentifiers", + identifier=identifier, + promoted=("supported_input_types", "supported_output_types"), + extra={ + "converter_target_hash": self.insert_target(converter_target) if converter_target else None, + "sub_converter_hash": self.insert_converter(sub_converter) if sub_converter else None, + }, + ) + + def insert_scenario(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scenario graph. + + Returns: + str | None: The stored hash when successful. + """ + objective_target = self._child(identifier, "objective_target") + objective_scorer = self._child(identifier, "objective_scorer") + return self._insert_identifier( + table="ScenarioIdentifiers", + identifier=identifier, + promoted=("version", "techniques", "datasets"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + + def insert_atomic_attack(self, identifier: dict[str, Any]) -> str | None: + """ + Insert an atomic attack graph. + + Returns: + str | None: The stored hash when successful. + """ + attack_technique = self._child(identifier, "attack_technique") + seeds = self._children(identifier, "seed_identifiers") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AtomicAttackIdentifiers", + identifier=identifier, + extra={ + "attack_technique_identifier_hash": ( + self._insert_attack_technique(attack_technique) if attack_technique else None + ) + }, + ) + if identifier_hash: + self._insert_edges( + table="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack_technique(self, identifier: dict[str, Any]) -> str | None: + attack = self._child(identifier, "attack") + seeds = self._children(identifier, "technique_seeds") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AttackTechniqueIdentifiers", + identifier=identifier, + extra={"attack_identifier_hash": self._insert_attack(attack) if attack else None}, + ) + if identifier_hash: + self._insert_edges( + table="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack(self, identifier: dict[str, Any]) -> str | None: + objective_target = self._child(identifier, "objective_target") + adversarial_chat = self._child(identifier, "adversarial_chat") + objective_scorer = self._child(identifier, "objective_scorer") + request_hashes = [ + value for item in self._children(identifier, "request_converters") if (value := self.insert_converter(item)) + ] + response_hashes = [ + value + for item in self._children(identifier, "response_converters") + if (value := self.insert_converter(item)) + ] + identifier_hash = self._insert_identifier( + table="AttackIdentifiers", + identifier=identifier, + promoted=("adversarial_system_prompt", "adversarial_seed_prompt"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "adversarial_chat_hash": self.insert_target(adversarial_chat) if adversarial_chat else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + if identifier_hash: + self._insert_edges( + table="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=request_hashes, + ) + self._insert_edges( + table="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=response_hashes, + ) + return identifier_hash + + def _insert_seed(self, identifier: dict[str, Any]) -> str | None: + return self._insert_identifier( + table="SeedIdentifiers", + identifier=identifier, + promoted=("value", "value_sha256", "data_type", "dataset_name", "is_general_technique"), + ) + + def _insert_identifier( + self, + *, + table: str, + identifier: dict[str, Any], + promoted: Sequence[str] = (), + extra: dict[str, Any] | None = None, + ) -> str | None: + identifier_hash = identifier.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64 or table not in self._hashes: + return None + if identifier_hash in self._hashes[table]: + return identifier_hash + values: dict[str, Any] = { + "hash": identifier_hash, + "class_name": identifier.get("class_name"), + "class_module": identifier.get("class_module"), + "identifier_json": json.dumps(identifier, sort_keys=True), + "pyrit_version": identifier.get("pyrit_version"), + } + values.update({name: self._json_value(identifier.get(name)) for name in promoted}) + values.update(extra or {}) + columns = list(values) + statement = sa.text( + f'INSERT INTO "{table}" ({", ".join(columns)}) ' + f'VALUES ({", ".join(f":{column}" for column in columns)})' + ) + self._bind.execute(statement, values) + self._hashes[table].add(identifier_hash) + return identifier_hash + + def _insert_edges( + self, + *, + table: str, + parent_column: str, + parent_hash: str, + child_column: str, + child_hashes: Sequence[str], + ) -> None: + statement = sa.text( + f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' + f'VALUES (:parent_hash, :position, :child_hash)' + ) + for position, child_hash in enumerate(child_hashes): + self._bind.execute( + statement, + {"parent_hash": parent_hash, "position": position, "child_hash": child_hash}, + ) + + @staticmethod + def _child( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> dict[str, Any] | None: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, dict): + return load_identifier(value) + return None + + @staticmethod + def _children( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> list[dict[str, Any]]: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, list): + return [child for item in value if (child := load_identifier(item)) is not None] + return [] + + @staticmethod + def _json_value(value: Any) -> Any: + return json.dumps(value) if isinstance(value, (list, dict)) else value diff --git a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py index 669a64716c..95c75c8ed4 100644 --- a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py +++ b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py @@ -11,13 +11,18 @@ from __future__ import annotations -import json import logging from collections.abc import Sequence # noqa: TC003 import sqlalchemy as sa from alembic import op +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier, + run_best_effort_backfill, +) + revision: str = "a6c8e0f2b4d6" down_revision: str | None = "e5f7a9c1b3d2" branch_labels: str | Sequence[str] | None = None @@ -31,9 +36,9 @@ def upgrade() -> None: op.create_table( "ScorerIdentifiers", sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=False), - sa.Column("class_module", sa.String(), nullable=False), - sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), sa.Column("scorer_type", sa.String(), nullable=True), sa.Column("score_aggregator", sa.String(), nullable=True), sa.Column("prompt_target_hash", sa.String(64), nullable=True), @@ -64,7 +69,8 @@ def upgrade() -> None: ["hash"], ) - _backfill_scorer_identifiers() + bind = op.get_bind() + run_best_effort_backfill(bind=bind, name="ScorerIdentifiers", backfill=_backfill_scorer_identifiers) def downgrade() -> None: @@ -77,107 +83,22 @@ def downgrade() -> None: def _backfill_scorer_identifiers() -> None: """Backfill scorer rows and score foreign keys from the retained JSON column.""" - from pyrit.models import ScorerIdentifier, TargetIdentifier - bind = op.get_bind() score_rows = bind.execute( sa.text('SELECT id, scorer_class_identifier FROM "ScoreEntries" WHERE scorer_class_identifier IS NOT NULL') ).fetchall() - scorer_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScorerIdentifiers"')).fetchall()} - target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} - - target_insert = sa.text( - 'INSERT INTO "TargetIdentifiers" ' - "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " - "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " - ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " - ":supported_auth_modes, :pyrit_version)" - ) - target_edge_insert = sa.text( - 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) - scorer_insert = sa.text( - 'INSERT INTO "ScorerIdentifiers" ' - "(hash, class_name, class_module, identifier_json, scorer_type, score_aggregator, " - "prompt_target_hash, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :scorer_type, " - ":score_aggregator, :prompt_target_hash, :pyrit_version)" - ) - scorer_edge_insert = sa.text( - 'INSERT INTO "ScorerIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') - - def _insert_target(identifier: TargetIdentifier) -> None: - if identifier.hash in target_hashes: - return - for child in identifier.targets: - _insert_target(child) - bind.execute( - target_insert, - { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "endpoint": identifier.endpoint, - "model_name": identifier.model_name, - "underlying_model_name": identifier.underlying_model_name, - "temperature": identifier.temperature, - "top_p": identifier.top_p, - "max_requests_per_minute": identifier.max_requests_per_minute, - "supported_auth_modes": ( - json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None - ), - "pyrit_version": identifier.pyrit_version, - }, - ) - target_hashes.add(identifier.hash) - for position, child in enumerate(identifier.targets): - bind.execute( - target_edge_insert, - {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, - ) - - def _insert_scorer(identifier: ScorerIdentifier) -> None: - if identifier.hash in scorer_hashes: - return - if identifier.prompt_target is not None: - _insert_target(identifier.prompt_target) - for child in identifier.sub_scorers: - _insert_scorer(child) - bind.execute( - scorer_insert, - { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "scorer_type": identifier.scorer_type, - "score_aggregator": identifier.score_aggregator, - "prompt_target_hash": identifier.prompt_target.hash if identifier.prompt_target is not None else None, - "pyrit_version": identifier.pyrit_version, - }, - ) - scorer_hashes.add(identifier.hash) - for position, child in enumerate(identifier.sub_scorers): - bind.execute( - scorer_edge_insert, - {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, - ) - + inserter = IdentifierGraphInserter(bind=bind) skipped = 0 for score_id, raw_scorer in score_rows: - stored = json.loads(raw_scorer) if isinstance(raw_scorer, str) else raw_scorer - if not stored: + identifier = load_identifier(raw_scorer) + if identifier is None: + skipped += 1 continue try: - identifier = ScorerIdentifier.model_validate(stored) - _insert_scorer(identifier) - bind.execute(score_update, {"hash": identifier.hash, "id": score_id}) + identifier_hash = inserter.insert_scorer(identifier) + if identifier_hash: + bind.execute(score_update, {"hash": identifier_hash, "id": score_id}) except Exception: skipped += 1 logger.warning( diff --git a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py index 859ef211c0..72a766bf82 100644 --- a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py +++ b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py @@ -11,13 +11,18 @@ from __future__ import annotations -import json import logging from collections.abc import Sequence # noqa: TC003 import sqlalchemy as sa from alembic import op +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier, + run_best_effort_backfill, +) + revision: str = "b7d9f1a3c5e7" down_revision: str | None = "a6c8e0f2b4d6" branch_labels: str | Sequence[str] | None = None @@ -31,9 +36,9 @@ def upgrade() -> None: op.create_table( "ScenarioIdentifiers", sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=False), - sa.Column("class_module", sa.String(), nullable=False), - sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), sa.Column("version", sa.Integer(), nullable=True), sa.Column("techniques", sa.JSON(), nullable=True), sa.Column("datasets", sa.JSON(), nullable=True), @@ -60,7 +65,8 @@ def upgrade() -> None: ["hash"], ) - _backfill_scenario_identifiers() + bind = op.get_bind() + run_best_effort_backfill(bind=bind, name="ScenarioIdentifiers", backfill=_backfill_scenario_identifiers) def downgrade() -> None: @@ -72,138 +78,22 @@ def downgrade() -> None: def _backfill_scenario_identifiers() -> None: """Backfill scenario rows and result foreign keys from the retained JSON column.""" - from pyrit.models import ScenarioIdentifier, ScorerIdentifier, TargetIdentifier - bind = op.get_bind() result_rows = bind.execute( sa.text('SELECT id, scenario_identifier FROM "ScenarioResultEntries" WHERE scenario_identifier IS NOT NULL') ).fetchall() - existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScenarioIdentifiers"')).fetchall()} - target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} - scorer_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ScorerIdentifiers"')).fetchall()} - target_insert = sa.text( - 'INSERT INTO "TargetIdentifiers" ' - "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " - "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " - ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " - ":supported_auth_modes, :pyrit_version)" - ) - target_edge_insert = sa.text( - 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) - scorer_insert = sa.text( - 'INSERT INTO "ScorerIdentifiers" ' - "(hash, class_name, class_module, identifier_json, scorer_type, score_aggregator, " - "prompt_target_hash, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :scorer_type, " - ":score_aggregator, :prompt_target_hash, :pyrit_version)" - ) - scorer_edge_insert = sa.text( - 'INSERT INTO "ScorerIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) - insert_stmt = sa.text( - 'INSERT INTO "ScenarioIdentifiers" ' - "(hash, class_name, class_module, identifier_json, version, techniques, datasets, " - "objective_target_hash, objective_scorer_hash, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :version, :techniques, :datasets, " - ":objective_target_hash, :objective_scorer_hash, :pyrit_version)" - ) update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') - - def _insert_target(identifier: TargetIdentifier) -> None: - if identifier.hash in target_hashes: - return - for child in identifier.targets: - _insert_target(child) - bind.execute( - target_insert, - { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "endpoint": identifier.endpoint, - "model_name": identifier.model_name, - "underlying_model_name": identifier.underlying_model_name, - "temperature": identifier.temperature, - "top_p": identifier.top_p, - "max_requests_per_minute": identifier.max_requests_per_minute, - "supported_auth_modes": ( - json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None - ), - "pyrit_version": identifier.pyrit_version, - }, - ) - target_hashes.add(identifier.hash) - for position, child in enumerate(identifier.targets): - bind.execute( - target_edge_insert, - {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, - ) - - def _insert_scorer(identifier: ScorerIdentifier) -> None: - if identifier.hash in scorer_hashes: - return - if identifier.prompt_target is not None: - _insert_target(identifier.prompt_target) - for child in identifier.sub_scorers: - _insert_scorer(child) - bind.execute( - scorer_insert, - { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "scorer_type": identifier.scorer_type, - "score_aggregator": identifier.score_aggregator, - "prompt_target_hash": identifier.prompt_target.hash if identifier.prompt_target is not None else None, - "pyrit_version": identifier.pyrit_version, - }, - ) - scorer_hashes.add(identifier.hash) - for position, child in enumerate(identifier.sub_scorers): - bind.execute( - scorer_edge_insert, - {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, - ) - + inserter = IdentifierGraphInserter(bind=bind) skipped = 0 for result_id, raw_scenario in result_rows: - stored = json.loads(raw_scenario) if isinstance(raw_scenario, str) else raw_scenario - if not stored: + identifier = load_identifier(raw_scenario) + if identifier is None: + skipped += 1 continue try: - identifier = ScenarioIdentifier.model_validate(stored) - if identifier.objective_target is not None: - _insert_target(identifier.objective_target) - if identifier.objective_scorer is not None: - _insert_scorer(identifier.objective_scorer) - if identifier.hash not in existing_hashes: - bind.execute( - insert_stmt, - { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "version": identifier.version, - "techniques": json.dumps(identifier.techniques) if identifier.techniques is not None else None, - "datasets": json.dumps(identifier.datasets) if identifier.datasets is not None else None, - "objective_target_hash": ( - identifier.objective_target.hash if identifier.objective_target is not None else None - ), - "objective_scorer_hash": ( - identifier.objective_scorer.hash if identifier.objective_scorer is not None else None - ), - "pyrit_version": identifier.pyrit_version, - }, - ) - existing_hashes.add(identifier.hash) - bind.execute(update_stmt, {"hash": identifier.hash, "id": result_id}) + identifier_hash = inserter.insert_scenario(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) except Exception: skipped += 1 logger.warning( diff --git a/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py b/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py index de0b529631..b114aea949 100644 --- a/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py +++ b/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py @@ -11,7 +11,6 @@ from __future__ import annotations -import json import logging import uuid from collections.abc import Sequence # noqa: TC003 @@ -22,6 +21,12 @@ from sqlalchemy.dialects.sqlite import CHAR from sqlalchemy.types import TypeDecorator, Uuid +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier_list, + run_best_effort_backfill, +) + if TYPE_CHECKING: from sqlalchemy.engine import Dialect @@ -61,9 +66,9 @@ def upgrade() -> None: op.create_table( "ConverterIdentifiers", sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=False), - sa.Column("class_module", sa.String(), nullable=False), - sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), sa.Column("supported_input_types", sa.JSON(), nullable=True), sa.Column("supported_output_types", sa.JSON(), nullable=True), sa.Column("converter_target_hash", sa.String(64), nullable=True), @@ -97,7 +102,8 @@ def upgrade() -> None: ), sa.PrimaryKeyConstraint("prompt_memory_entry_id", "position"), ) - _backfill_converter_identifiers() + bind = op.get_bind() + run_best_effort_backfill(bind=bind, name="ConverterIdentifiers", backfill=_backfill_converter_identifiers) def downgrade() -> None: @@ -108,8 +114,6 @@ def downgrade() -> None: def _backfill_converter_identifiers() -> None: """Materialize converter graphs and prompt associations from retained JSON.""" - from pyrit.models import ComponentIdentifier, ConverterIdentifier, TargetIdentifier - bind = op.get_bind() prompt_rows = bind.execute( sa.text( @@ -117,98 +121,34 @@ def _backfill_converter_identifiers() -> None: "WHERE converter_identifiers IS NOT NULL" ) ).fetchall() - converter_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "ConverterIdentifiers"'))} - target_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"'))} - - target_insert = sa.text( - 'INSERT INTO "TargetIdentifiers" ' - "(hash, class_name, class_module, identifier_json, endpoint, model_name, underlying_model_name, " - "temperature, top_p, max_requests_per_minute, supported_auth_modes, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " - ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " - ":supported_auth_modes, :pyrit_version)" - ) - target_edge_insert = sa.text( - 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) - converter_insert = sa.text( - 'INSERT INTO "ConverterIdentifiers" ' - "(hash, class_name, class_module, identifier_json, supported_input_types, supported_output_types, " - "converter_target_hash, sub_converter_hash, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :supported_input_types, " - ":supported_output_types, :converter_target_hash, :sub_converter_hash, :pyrit_version)" - ) link_insert = sa.text( 'INSERT INTO "PromptConverterIdentifiers" ' "(prompt_memory_entry_id, position, converter_identifier_hash) " "VALUES (:prompt_memory_entry_id, :position, :converter_identifier_hash)" ) - def _insert_target(identifier: TargetIdentifier) -> None: - if identifier.hash in target_hashes: - return - for child in identifier.targets: - _insert_target(child) - bind.execute(target_insert, _identifier_values(identifier)) - for position, child in enumerate(identifier.targets): - bind.execute( - target_edge_insert, - {"parent_hash": identifier.hash, "position": position, "child_hash": child.hash}, - ) - target_hashes.add(identifier.hash) - - def _insert_converter(identifier: ConverterIdentifier) -> None: - if identifier.hash in converter_hashes: - return - if identifier.converter_target is not None: - _insert_target(identifier.converter_target) - if identifier.sub_converter is not None: - _insert_converter(identifier.sub_converter) - bind.execute(converter_insert, _identifier_values(identifier)) - converter_hashes.add(identifier.hash) - + inserter = IdentifierGraphInserter(bind=bind) skipped = 0 for prompt_id, stored_identifiers, pyrit_version in prompt_rows: try: - values = json.loads(stored_identifiers) if isinstance(stored_identifiers, str) else stored_identifiers - for position, stored_identifier in enumerate(values): - stored_identifier["pyrit_version"] = pyrit_version - identifier = ConverterIdentifier.from_component_identifier( - ComponentIdentifier.model_validate(stored_identifier) - ) - _insert_converter(identifier) - bind.execute( - link_insert, - { - "prompt_memory_entry_id": prompt_id, - "position": position, - "converter_identifier_hash": identifier.hash, - }, - ) - except (TypeError, ValueError, KeyError): + for position, identifier in enumerate(load_identifier_list(stored_identifiers)): + if identifier.get("pyrit_version") is None: + identifier = {**identifier, "pyrit_version": pyrit_version} + identifier_hash = inserter.insert_converter(identifier) + if identifier_hash: + bind.execute( + link_insert, + { + "prompt_memory_entry_id": prompt_id, + "position": position, + "converter_identifier_hash": identifier_hash, + }, + ) + except Exception: skipped += 1 - logger.warning(f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}") + logger.warning( + f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", + exc_info=True, + ) if skipped: logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") - - -def _identifier_values(identifier: Any) -> dict[str, Any]: - """Return common and promoted values for a normalized identifier insert.""" - promoted_values = { - name: json.dumps(value) if isinstance(value, (list, dict)) else value - for name, value in identifier.promoted_scalar_values().items() - } - return { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump()), - "pyrit_version": identifier.pyrit_version, - **promoted_values, - **{ - f"{field_name}_hash": child.hash if child is not None else None - for field_name in identifier.promoted_child_field_names() - if not isinstance((child := getattr(identifier, field_name)), list) - }, - } diff --git a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py index 426e8295d7..0606d9856a 100644 --- a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py +++ b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py @@ -11,7 +11,6 @@ from __future__ import annotations -import json import logging from collections.abc import Sequence # noqa: TC003 from typing import Any @@ -19,6 +18,12 @@ import sqlalchemy as sa from alembic import op +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier, + run_best_effort_backfill, +) + revision: str = "d9f2a4b6c8e0" down_revision: str | None = "c8e1f3a5b7d9" branch_labels: str | Sequence[str] | None = None @@ -38,7 +43,8 @@ def upgrade() -> None: ["atomic_attack_identifier_hash"], ["hash"], ) - _backfill_attack_identifiers() + bind = op.get_bind() + run_best_effort_backfill(bind=bind, name="AttackIdentifiers", backfill=_backfill_attack_identifiers) def downgrade() -> None: @@ -125,9 +131,9 @@ def _create_identifier_tables() -> None: def _common_columns() -> tuple[sa.Column[Any], ...]: return ( sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=False), - sa.Column("class_module", sa.String(), nullable=False), - sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), sa.Column("pyrit_version", sa.String(), nullable=True), ) @@ -153,8 +159,6 @@ def _create_ordered_edge_table( def _backfill_attack_identifiers() -> None: """Backfill attack identifier graphs and result links from retained JSON.""" - from pyrit.models import AtomicAttackIdentifier - bind = op.get_bind() result_rows = bind.execute( sa.text( @@ -162,18 +166,21 @@ def _backfill_attack_identifiers() -> None: "WHERE atomic_attack_identifier IS NOT NULL" ) ).fetchall() - inserter = _AttackIdentifierGraphInserter(bind=bind) + inserter = IdentifierGraphInserter(bind=bind) update_stmt = sa.text( 'UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id' ) skipped = 0 for result_id, raw_identifier in result_rows: + identifier = load_identifier(raw_identifier) + if identifier is None: + skipped += 1 + continue try: - stored = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier - identifier = AtomicAttackIdentifier.model_validate(stored) - inserter.insert_atomic_attack(identifier) - bind.execute(update_stmt, {"hash": identifier.hash, "id": result_id}) + identifier_hash = inserter.insert_atomic_attack(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) except Exception: skipped += 1 logger.warning( @@ -182,220 +189,3 @@ def _backfill_attack_identifiers() -> None: ) if skipped: logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") - - -class _AttackIdentifierGraphInserter: - """Insert a normalized attack identifier graph during migration backfill.""" - - def __init__(self, *, bind: Any) -> None: - self._bind = bind - self._hashes = { - table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) - for table in ( - "TargetIdentifiers", - "ScorerIdentifiers", - "ConverterIdentifiers", - "SeedIdentifiers", - "AttackIdentifiers", - "AttackTechniqueIdentifiers", - "AtomicAttackIdentifiers", - ) - } - - def insert_atomic_attack(self, identifier: Any) -> None: - if identifier.hash in self._hashes["AtomicAttackIdentifiers"]: - return - if identifier.attack_technique is not None: - self._insert_attack_technique(identifier.attack_technique) - for seed in identifier.seed_identifiers: - self._insert_seed(seed) - self._insert_identifier( - table="AtomicAttackIdentifiers", - identifier=identifier, - extra={ - "attack_technique_identifier_hash": ( - identifier.attack_technique.hash if identifier.attack_technique is not None else None - ) - }, - ) - self._insert_edges( - table="AtomicAttackSeedIdentifiers", - parent_column="atomic_attack_identifier_hash", - parent_hash=identifier.hash, - child_column="seed_identifier_hash", - children=identifier.seed_identifiers, - ) - - def _insert_attack_technique(self, identifier: Any) -> None: - if identifier.hash in self._hashes["AttackTechniqueIdentifiers"]: - return - if identifier.attack is not None: - self._insert_attack(identifier.attack) - for seed in identifier.technique_seeds: - self._insert_seed(seed) - self._insert_identifier( - table="AttackTechniqueIdentifiers", - identifier=identifier, - extra={"attack_identifier_hash": identifier.attack.hash if identifier.attack is not None else None}, - ) - self._insert_edges( - table="AttackTechniqueSeedIdentifiers", - parent_column="attack_technique_identifier_hash", - parent_hash=identifier.hash, - child_column="seed_identifier_hash", - children=identifier.technique_seeds, - ) - - def _insert_attack(self, identifier: Any) -> None: - if identifier.hash in self._hashes["AttackIdentifiers"]: - return - for target in (identifier.objective_target, identifier.adversarial_chat): - if target is not None: - self._insert_target(target) - if identifier.objective_scorer is not None: - self._insert_scorer(identifier.objective_scorer) - for converter in [*identifier.request_converters, *identifier.response_converters]: - self._insert_converter(converter) - self._insert_identifier( - table="AttackIdentifiers", - identifier=identifier, - extra={ - "adversarial_system_prompt": identifier.adversarial_system_prompt, - "adversarial_seed_prompt": identifier.adversarial_seed_prompt, - "objective_target_hash": ( - identifier.objective_target.hash if identifier.objective_target is not None else None - ), - "adversarial_chat_hash": ( - identifier.adversarial_chat.hash if identifier.adversarial_chat is not None else None - ), - "objective_scorer_hash": ( - identifier.objective_scorer.hash if identifier.objective_scorer is not None else None - ), - }, - ) - self._insert_edges( - table="AttackRequestConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_hash=identifier.hash, - child_column="converter_identifier_hash", - children=identifier.request_converters, - ) - self._insert_edges( - table="AttackResponseConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_hash=identifier.hash, - child_column="converter_identifier_hash", - children=identifier.response_converters, - ) - - def _insert_seed(self, identifier: Any) -> None: - if identifier.hash in self._hashes["SeedIdentifiers"]: - return - self._insert_identifier( - table="SeedIdentifiers", - identifier=identifier, - extra=identifier.promoted_scalar_values(), - ) - - def _insert_target(self, identifier: Any) -> None: - if identifier.hash in self._hashes["TargetIdentifiers"]: - return - for child in identifier.targets: - self._insert_target(child) - self._insert_identifier( - table="TargetIdentifiers", - identifier=identifier, - extra=identifier.promoted_scalar_values(), - ) - self._insert_edges( - table="TargetIdentifierChildren", - parent_column="parent_hash", - parent_hash=identifier.hash, - child_column="child_hash", - children=identifier.targets, - ) - - def _insert_scorer(self, identifier: Any) -> None: - if identifier.hash in self._hashes["ScorerIdentifiers"]: - return - if identifier.prompt_target is not None: - self._insert_target(identifier.prompt_target) - for child in identifier.sub_scorers: - self._insert_scorer(child) - self._insert_identifier( - table="ScorerIdentifiers", - identifier=identifier, - extra={ - **identifier.promoted_scalar_values(), - "prompt_target_hash": ( - identifier.prompt_target.hash if identifier.prompt_target is not None else None - ), - }, - ) - self._insert_edges( - table="ScorerIdentifierChildren", - parent_column="parent_hash", - parent_hash=identifier.hash, - child_column="child_hash", - children=identifier.sub_scorers, - ) - - def _insert_converter(self, identifier: Any) -> None: - if identifier.hash in self._hashes["ConverterIdentifiers"]: - return - if identifier.converter_target is not None: - self._insert_target(identifier.converter_target) - if identifier.sub_converter is not None: - self._insert_converter(identifier.sub_converter) - self._insert_identifier( - table="ConverterIdentifiers", - identifier=identifier, - extra={ - **identifier.promoted_scalar_values(), - "converter_target_hash": ( - identifier.converter_target.hash if identifier.converter_target is not None else None - ), - "sub_converter_hash": ( - identifier.sub_converter.hash if identifier.sub_converter is not None else None - ), - }, - ) - - def _insert_identifier(self, *, table: str, identifier: Any, extra: dict[str, Any]) -> None: - columns = ["hash", "class_name", "class_module", "identifier_json", *extra, "pyrit_version"] - placeholders = [f":{column}" for column in columns] - values = { - "hash": identifier.hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "pyrit_version": identifier.pyrit_version, - **{ - name: json.dumps(value) if isinstance(value, (list, dict)) else value - for name, value in extra.items() - }, - } - self._bind.execute( - sa.text(f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(placeholders)})'), - values, - ) - self._hashes[table].add(identifier.hash) - - def _insert_edges( - self, - *, - table: str, - parent_column: str, - parent_hash: str, - child_column: str, - children: Sequence[Any], - ) -> None: - statement = sa.text( - f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' - f'VALUES (:parent_hash, :position, :child_hash)' - ) - for position, child in enumerate(children): - self._bind.execute( - statement, - {"parent_hash": parent_hash, "position": position, "child_hash": child.hash}, - ) diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py index f64af89ef6..4ed290f2ed 100644 --- a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py @@ -21,13 +21,18 @@ from __future__ import annotations -import json import logging from collections.abc import Sequence # noqa: TC003 import sqlalchemy as sa from alembic import op +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier, + run_best_effort_backfill, +) + # revision identifiers, used by Alembic. revision: str = "e5f7a9c1b3d2" down_revision: str | None = "d4e6f8a0b2c4" @@ -43,9 +48,9 @@ def upgrade() -> None: op.create_table( "TargetIdentifiers", sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=False), - sa.Column("class_module", sa.String(), nullable=False), - sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), sa.Column("endpoint", sa.String(), nullable=True), sa.Column("model_name", sa.String(), nullable=True), sa.Column("underlying_model_name", sa.String(), nullable=True), @@ -86,7 +91,12 @@ def upgrade() -> None: ["hash"], ) - _backfill_target_identifiers() + bind = op.get_bind() + run_best_effort_backfill( + bind=bind, + name="TargetIdentifiers", + backfill=_backfill_target_identifiers, + ) def downgrade() -> None: @@ -112,92 +122,28 @@ def _backfill_target_identifiers() -> None: are not re-inserted. Rows whose stored target cannot be reconstructed are logged and skipped rather than aborting the upgrade. """ - from pyrit.models import TargetIdentifier - bind = op.get_bind() rows = bind.execute( sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') ).fetchall() - existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} - - insert_stmt = sa.text( - 'INSERT INTO "TargetIdentifiers" ' - "(hash, class_name, class_module, identifier_json, endpoint, model_name, " - "underlying_model_name, temperature, top_p, max_requests_per_minute, " - "supported_auth_modes, pyrit_version) " - "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " - ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " - ":supported_auth_modes, :pyrit_version)" - ) - edge_stmt = sa.text( - 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' - "VALUES (:parent_hash, :position, :child_hash)" - ) update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') - - def _insert_target(identifier: TargetIdentifier) -> int: - """ - Insert ``identifier`` and its descendants if absent; record child edges. - - Returns: - int: The number of new ``TargetIdentifiers`` rows inserted (self + descendants). - """ - target_hash = identifier.hash - if target_hash in existing_hashes: - return 0 - # Children first so the parent's edge foreign keys resolve. - inserted = 0 - for child in identifier.targets: - inserted += _insert_target(child) - bind.execute( - insert_stmt, - { - "hash": target_hash, - "class_name": identifier.class_name, - "class_module": identifier.class_module, - "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), - "endpoint": identifier.endpoint, - "model_name": identifier.model_name, - "underlying_model_name": identifier.underlying_model_name, - "temperature": identifier.temperature, - "top_p": identifier.top_p, - "max_requests_per_minute": identifier.max_requests_per_minute, - "supported_auth_modes": ( - json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None - ), - "pyrit_version": identifier.pyrit_version, - }, - ) - existing_hashes.add(target_hash) - inserted += 1 - for position, child in enumerate(identifier.targets): - bind.execute(edge_stmt, {"parent_hash": target_hash, "position": position, "child_hash": child.hash}) - return inserted - - inserted = 0 + inserter = IdentifierGraphInserter(bind=bind) linked = 0 skipped = 0 for conversation_id, raw_target in rows: - stored = json.loads(raw_target) if isinstance(raw_target, str) else raw_target - if not stored: + identifier = load_identifier(raw_target) + if identifier is None: + skipped += 1 continue try: - identifier = TargetIdentifier.model_validate(stored) + identifier_hash = inserter.insert_target(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "cid": conversation_id}) + linked += 1 except Exception: skipped += 1 - logger.warning( - f"TargetIdentifiers backfill: could not reconstruct target for " - f"conversation_id {conversation_id!r}; skipping." - ) - continue - - inserted += _insert_target(identifier) - bind.execute(update_stmt, {"hash": identifier.hash, "cid": conversation_id}) - linked += 1 + logger.warning(f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", exc_info=True) - if inserted or linked or skipped: - logger.info( - f"TargetIdentifiers backfill: inserted {inserted} identifier row(s); " - f"linked {linked} conversation(s); skipped {skipped}." - ) + if linked or skipped: + logger.info(f"TargetIdentifiers backfill linked {linked} conversation(s); skipped {skipped}.") diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index cdd7948985..5162275468 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -445,12 +445,12 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): Abstract base for tables that persist a ``ComponentIdentifier`` projection. Mirrors the identifier class hierarchy: concrete identifier tables inherit the - shared, always-populated columns the way ``TargetIdentifier`` inherits + shared projection columns the way ``TargetIdentifier`` inherits ``ComponentIdentifier``. The content ``hash`` is the natural, dedupable primary - key, and ``identifier_json`` holds the full flat dump for lossless - reconstruction (children stay inline). Rows are immutable — the same content - always maps to the same hash, so a given identifier reused across rows is - stored once. + key. Runtime writes populate the descriptive fields and full ``identifier_json``; + they remain nullable so best-effort migration backfills can preserve partial + legacy identifiers. Rows are immutable — the same content always maps to the + same hash, so a given identifier reused across rows is stored once. Subclasses declare their promoted query columns and implement the ``DomainBackedEntry.from_domain_model`` seam to map their strongly-typed identifier @@ -469,10 +469,10 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): #: Content-addressed identity — the same value as ``ComponentIdentifier.hash``. #: SHA256 hex digest is 64 chars; bounded for SQL Server key/index compatibility. hash: Mapped[str] = mapped_column(String(64), primary_key=True) - class_name: Mapped[str] = mapped_column(String, nullable=False) - class_module: Mapped[str] = mapped_column(String, nullable=False) + class_name: Mapped[str | None] = mapped_column(String, nullable=True) + class_module: Mapped[str | None] = mapped_column(String, nullable=True) #: Full flat ``model_dump()`` of the identifier. Source of truth on reload. - identifier_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) + identifier_json: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) #: Version that first wrote this content-addressed row. Nullable for backwards #: compatibility with existing databases. pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 17988a8f4f..47e4d36843 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -814,6 +814,124 @@ def test_conversations_migration_downgrade_restores_columns(): # ============================================================================= +def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json(): + """Malformed retained identifiers do not block upgrades and leave nullable links unset.""" + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-best-effort.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "d4e6f8a0b2c4") + connection.execute( + text('INSERT INTO "Conversations" (conversation_id, target_identifier) VALUES (:id, :value)'), + {"id": "malformed-conversation", "value": "not-json"}, + ) + command.upgrade(config, "e5f7a9c1b3d2") + assert connection.execute( + text( + 'SELECT target_identifier_hash FROM "Conversations" ' + "WHERE conversation_id = 'malformed-conversation'" + ) + ).scalar_one() is None + + score_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "ScoreEntries" ' + "(id, score_value, score_type, score_metadata, scorer_class_identifier, timestamp) " + "VALUES (:id, 'True', 'true_false', '{}', 'not-json', '2026-07-13')" + ), + {"id": score_id}, + ) + command.upgrade(config, "a6c8e0f2b4d6") + assert connection.execute( + text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), + {"id": score_id}, + ).scalar_one() is None + + result_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "ScenarioResultEntries" ' + "(id, scenario_name, scenario_version, pyrit_version, scenario_identifier, " + "objective_target_identifier, scenario_run_state, attack_results_json, number_tries, " + "completion_time, timestamp) VALUES (:id, 'Scenario', 1, '0.10.0', 'not-json', '{}', " + "'COMPLETED', '{}', 1, '2026-07-13', '2026-07-13')" + ), + {"id": result_id}, + ) + command.upgrade(config, "b7d9f1a3c5e7") + assert connection.execute( + text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() is None + + prompt_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "PromptMemoryEntries" ' + "(id, role, conversation_id, sequence, timestamp, labels, prompt_metadata, " + "converter_identifiers, original_value_data_type, original_value, converted_value_data_type, " + "original_prompt_id) VALUES (:id, 'user', 'conversation', 0, '2026-07-13', '{}', '{}', " + "'not-json', 'text', 'prompt', 'text', :id)" + ), + {"id": prompt_id}, + ) + command.upgrade(config, "c8e1f3a5b7d9") + assert connection.execute(text('SELECT COUNT(*) FROM "PromptConverterIdentifiers"')).scalar_one() == 0 + + attack_result_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, atomic_attack_identifier, executed_turns, execution_time_ms, " + "outcome, timestamp) VALUES (:id, 'conversation', 'objective', 'not-json', 0, 0, " + "'undetermined', '2026-07-13')" + ), + {"id": attack_result_id}, + ) + command.upgrade(config, "d9f2a4b6c8e0") + assert connection.execute( + text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), + {"id": attack_result_id}, + ).scalar_one() is None + + for table_name in ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ScenarioIdentifiers", + "ConverterIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ): + columns = {column["name"]: column for column in inspect(connection).get_columns(table_name)} + assert columns["hash"]["nullable"] is False + assert columns["class_name"]["nullable"] is True + assert columns["class_module"]["nullable"] is True + assert columns["identifier_json"]["nullable"] is True + finally: + engine.dispose() + + +def test_identifier_migrations_do_not_import_domain_models(): + """Frozen identifier migrations operate on retained JSON rather than current domain models.""" + versions_dir = Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic" / "versions" + revision_names = ( + "e5f7a9c1b3d2_add_target_identifiers_table.py", + "a6c8e0f2b4d6_add_scorer_identifiers_table.py", + "b7d9f1a3c5e7_add_scenario_identifiers_table.py", + "c8e1f3a5b7d9_add_converter_identifiers_table.py", + "d9f2a4b6c8e0_add_attack_identifiers_tables.py", + ) + + for revision_name in revision_names: + source = (versions_dir / revision_name).read_text(encoding="utf-8") + assert "from pyrit.models" not in source + assert "import pyrit.models" not in source + + def test_scorer_identifier_migration_backfills_graph_and_score_link(): """Existing scorer JSON is materialized and linked without changing its content identity.""" from pyrit.models import ComponentIdentifier From a81082bb38a263cec1a1d9a77baa36fb31bade54 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 10:06:45 -0700 Subject: [PATCH 12/24] precommit --- pyrit/memory/alembic/identifier_backfill.py | 5 +- ...2a4b6c8e0_add_attack_identifiers_tables.py | 7 +- pyrit/memory/memory_models.py | 4 +- tests/unit/memory/test_memory_models.py | 4 +- tests/unit/memory/test_migration.py | 64 +++++++++++-------- 5 files changed, 44 insertions(+), 40 deletions(-) diff --git a/pyrit/memory/alembic/identifier_backfill.py b/pyrit/memory/alembic/identifier_backfill.py index a883e7746c..3a3b22c522 100644 --- a/pyrit/memory/alembic/identifier_backfill.py +++ b/pyrit/memory/alembic/identifier_backfill.py @@ -300,8 +300,7 @@ def _insert_identifier( values.update(extra or {}) columns = list(values) statement = sa.text( - f'INSERT INTO "{table}" ({", ".join(columns)}) ' - f'VALUES ({", ".join(f":{column}" for column in columns)})' + f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(f":{column}" for column in columns)})' ) self._bind.execute(statement, values) self._hashes[table].add(identifier_hash) @@ -318,7 +317,7 @@ def _insert_edges( ) -> None: statement = sa.text( f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' - f'VALUES (:parent_hash, :position, :child_hash)' + f"VALUES (:parent_hash, :position, :child_hash)" ) for position, child_hash in enumerate(child_hashes): self._bind.execute( diff --git a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py index 0606d9856a..7e7bde32df 100644 --- a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py +++ b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py @@ -162,14 +162,11 @@ def _backfill_attack_identifiers() -> None: bind = op.get_bind() result_rows = bind.execute( sa.text( - 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" ' - "WHERE atomic_attack_identifier IS NOT NULL" + 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" WHERE atomic_attack_identifier IS NOT NULL' ) ).fetchall() inserter = IdentifierGraphInserter(bind=bind) - update_stmt = sa.text( - 'UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id' - ) + update_stmt = sa.text('UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id') skipped = 0 for result_id, raw_identifier in result_rows: diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 5162275468..8aa30514ac 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1581,9 +1581,7 @@ def __init__(self, *, entry: AttackResult) -> None: AtomicAttackEvaluationIdentifier(atomic_attack_identifier).eval_hash ) entry.atomic_attack_identifier = atomic_attack_identifier - self.atomic_attack_identifier = ( - atomic_attack_identifier.model_dump() if atomic_attack_identifier else None - ) + self.atomic_attack_identifier = atomic_attack_identifier.model_dump() if atomic_attack_identifier else None self.atomic_attack_identifier_hash = atomic_attack_identifier.hash if atomic_attack_identifier else None self.objective_sha256 = to_sha256(entry.objective) diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 759d5d9424..52107984d0 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -294,8 +294,8 @@ def test_atomic_attack_identifier_graph_persists_with_result_link() -> None: memory = MagicMock(spec=MemoryInterface) memory.get_session.side_effect = lambda: Session(engine) - memory._persist_identifier.side_effect = ( - lambda *, session, identifier: MemoryInterface._persist_identifier(session=session, identifier=identifier) + memory._persist_identifier.side_effect = lambda *, session, identifier: MemoryInterface._persist_identifier( + session=session, identifier=identifier ) MemoryInterface.add_attack_results_to_memory(memory, attack_results=[result]) diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 47e4d36843..715fb96389 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -827,12 +827,15 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( {"id": "malformed-conversation", "value": "not-json"}, ) command.upgrade(config, "e5f7a9c1b3d2") - assert connection.execute( - text( - 'SELECT target_identifier_hash FROM "Conversations" ' - "WHERE conversation_id = 'malformed-conversation'" - ) - ).scalar_one() is None + assert ( + connection.execute( + text( + 'SELECT target_identifier_hash FROM "Conversations" ' + "WHERE conversation_id = 'malformed-conversation'" + ) + ).scalar_one() + is None + ) score_id = str(uuid.uuid4()) connection.execute( @@ -844,10 +847,13 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( {"id": score_id}, ) command.upgrade(config, "a6c8e0f2b4d6") - assert connection.execute( - text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), - {"id": score_id}, - ).scalar_one() is None + assert ( + connection.execute( + text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), + {"id": score_id}, + ).scalar_one() + is None + ) result_id = str(uuid.uuid4()) connection.execute( @@ -861,10 +867,13 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( {"id": result_id}, ) command.upgrade(config, "b7d9f1a3c5e7") - assert connection.execute( - text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), - {"id": result_id}, - ).scalar_one() is None + assert ( + connection.execute( + text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() + is None + ) prompt_id = str(uuid.uuid4()) connection.execute( @@ -891,10 +900,13 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( {"id": attack_result_id}, ) command.upgrade(config, "d9f2a4b6c8e0") - assert connection.execute( - text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), - {"id": attack_result_id}, - ).scalar_one() is None + assert ( + connection.execute( + text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), + {"id": attack_result_id}, + ).scalar_one() + is None + ) for table_name in ( "TargetIdentifiers", @@ -1246,18 +1258,16 @@ def test_attack_identifier_migration_backfills_graph_and_result_link(): text('SELECT hash, attack_identifier_hash FROM "AttackTechniqueIdentifiers"') ).one() attack_row = connection.execute( - text( - "SELECT hash, objective_target_hash, objective_scorer_hash " - 'FROM "AttackIdentifiers"' - ) + text('SELECT hash, objective_target_hash, objective_scorer_hash FROM "AttackIdentifiers"') ).one() seed_hashes = set(connection.execute(text('SELECT hash FROM "SeedIdentifiers"')).scalars()) - atomic_seed_hashes = connection.execute( - text( - "SELECT seed_identifier_hash FROM " - '"AtomicAttackSeedIdentifiers" ORDER BY position' + atomic_seed_hashes = ( + connection.execute( + text('SELECT seed_identifier_hash FROM "AtomicAttackSeedIdentifiers" ORDER BY position') ) - ).scalars().all() + .scalars() + .all() + ) request_converter_hash = connection.execute( text('SELECT converter_identifier_hash FROM "AttackRequestConverterIdentifiers"') ).scalar_one() From 571c492a255af69d0b844736280d9ee207ecdf9c Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 10:09:28 -0700 Subject: [PATCH 13/24] ruff --- pyrit/memory/memory_interface.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 24289a0d9b..acd7a604b8 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -424,6 +424,9 @@ def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece Args: message_pieces (Sequence[MessagePiece]): Persistable pieces (none flagged ``not_in_memory``), each carrying a non-empty ``conversation_id``. + + Raises: + SQLAlchemyError: If the message pieces or converter identifiers cannot be persisted. """ entries = [PromptMemoryEntry(entry=piece) for piece in message_pieces] with closing(self.get_session()) as session: From 14fa39aefa222320e7cb9bc9581fe38e209e4c21 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 10:34:17 -0700 Subject: [PATCH 14/24] merge scripts --- ...c8e0f2b4d6_add_scorer_identifiers_table.py | 110 ---- ...f1a3c5e7_add_scenario_identifiers_table.py | 105 ---- ...3a5b7d9_add_converter_identifiers_table.py | 154 ------ ...2a4b6c8e0_add_attack_identifiers_tables.py | 188 ------- .../e5f7a9c1b3d2_add_identifiers_tables.py | 519 ++++++++++++++++++ ...f7a9c1b3d2_add_target_identifiers_table.py | 149 ----- tests/unit/memory/test_migration.py | 92 ++-- 7 files changed, 559 insertions(+), 758 deletions(-) delete mode 100644 pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py delete mode 100644 pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py delete mode 100644 pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py delete mode 100644 pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py create mode 100644 pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py delete mode 100644 pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py deleted file mode 100644 index 95c75c8ed4..0000000000 --- a/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Persist scorer identifiers as content-addressed rows. - -Revision ID: a6c8e0f2b4d6 -Revises: e5f7a9c1b3d2 -Create Date: 2026-07-13 12:00:00.000000 -""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence # noqa: TC003 - -import sqlalchemy as sa -from alembic import op - -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier, - run_best_effort_backfill, -) - -revision: str = "a6c8e0f2b4d6" -down_revision: str | None = "e5f7a9c1b3d2" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -logger = logging.getLogger(__name__) - - -def upgrade() -> None: - """Apply this schema upgrade.""" - op.create_table( - "ScorerIdentifiers", - sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=True), - sa.Column("class_module", sa.String(), nullable=True), - sa.Column("identifier_json", sa.JSON(), nullable=True), - sa.Column("scorer_type", sa.String(), nullable=True), - sa.Column("score_aggregator", sa.String(), nullable=True), - sa.Column("prompt_target_hash", sa.String(64), nullable=True), - sa.Column("pyrit_version", sa.String(), nullable=True), - sa.ForeignKeyConstraint( - ["prompt_target_hash"], ["TargetIdentifiers.hash"], name="fk_scorer_identifiers_prompt_target_hash" - ), - ) - op.create_table( - "ScorerIdentifierChildren", - sa.Column("parent_hash", sa.String(64), nullable=False), - sa.Column("position", sa.Integer(), nullable=False), - sa.Column("child_hash", sa.String(64), nullable=False), - sa.PrimaryKeyConstraint("parent_hash", "position"), - sa.ForeignKeyConstraint( - ["parent_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_parent_hash" - ), - sa.ForeignKeyConstraint( - ["child_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_child_hash" - ), - ) - with op.batch_alter_table("ScoreEntries") as batch_op: - batch_op.add_column(sa.Column("scorer_identifier_hash", sa.String(64), nullable=True)) - batch_op.create_foreign_key( - "fk_score_entries_scorer_identifier_hash", - "ScorerIdentifiers", - ["scorer_identifier_hash"], - ["hash"], - ) - - bind = op.get_bind() - run_best_effort_backfill(bind=bind, name="ScorerIdentifiers", backfill=_backfill_scorer_identifiers) - - -def downgrade() -> None: - """Revert this schema upgrade.""" - with op.batch_alter_table("ScoreEntries") as batch_op: - batch_op.drop_column("scorer_identifier_hash") - op.drop_table("ScorerIdentifierChildren") - op.drop_table("ScorerIdentifiers") - - -def _backfill_scorer_identifiers() -> None: - """Backfill scorer rows and score foreign keys from the retained JSON column.""" - bind = op.get_bind() - score_rows = bind.execute( - sa.text('SELECT id, scorer_class_identifier FROM "ScoreEntries" WHERE scorer_class_identifier IS NOT NULL') - ).fetchall() - score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') - inserter = IdentifierGraphInserter(bind=bind) - skipped = 0 - for score_id, raw_scorer in score_rows: - identifier = load_identifier(raw_scorer) - if identifier is None: - skipped += 1 - continue - try: - identifier_hash = inserter.insert_scorer(identifier) - if identifier_hash: - bind.execute(score_update, {"hash": identifier_hash, "id": score_id}) - except Exception: - skipped += 1 - logger.warning( - f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", - exc_info=True, - ) - - if skipped: - logger.warning(f"ScorerIdentifiers backfill skipped {skipped} score row(s)") diff --git a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py deleted file mode 100644 index 72a766bf82..0000000000 --- a/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Persist scenario identifiers as content-addressed rows. - -Revision ID: b7d9f1a3c5e7 -Revises: a6c8e0f2b4d6 -Create Date: 2026-07-13 13:00:00.000000 -""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence # noqa: TC003 - -import sqlalchemy as sa -from alembic import op - -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier, - run_best_effort_backfill, -) - -revision: str = "b7d9f1a3c5e7" -down_revision: str | None = "a6c8e0f2b4d6" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -logger = logging.getLogger(__name__) - - -def upgrade() -> None: - """Apply this schema upgrade.""" - op.create_table( - "ScenarioIdentifiers", - sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=True), - sa.Column("class_module", sa.String(), nullable=True), - sa.Column("identifier_json", sa.JSON(), nullable=True), - sa.Column("version", sa.Integer(), nullable=True), - sa.Column("techniques", sa.JSON(), nullable=True), - sa.Column("datasets", sa.JSON(), nullable=True), - sa.Column("objective_target_hash", sa.String(64), nullable=True), - sa.Column("objective_scorer_hash", sa.String(64), nullable=True), - sa.Column("pyrit_version", sa.String(), nullable=True), - sa.ForeignKeyConstraint( - ["objective_target_hash"], - ["TargetIdentifiers.hash"], - name="fk_scenario_identifiers_objective_target_hash", - ), - sa.ForeignKeyConstraint( - ["objective_scorer_hash"], - ["ScorerIdentifiers.hash"], - name="fk_scenario_identifiers_objective_scorer_hash", - ), - ) - with op.batch_alter_table("ScenarioResultEntries") as batch_op: - batch_op.add_column(sa.Column("scenario_identifier_hash", sa.String(64), nullable=True)) - batch_op.create_foreign_key( - "fk_scenario_result_entries_scenario_identifier_hash", - "ScenarioIdentifiers", - ["scenario_identifier_hash"], - ["hash"], - ) - - bind = op.get_bind() - run_best_effort_backfill(bind=bind, name="ScenarioIdentifiers", backfill=_backfill_scenario_identifiers) - - -def downgrade() -> None: - """Revert this schema upgrade.""" - with op.batch_alter_table("ScenarioResultEntries") as batch_op: - batch_op.drop_column("scenario_identifier_hash") - op.drop_table("ScenarioIdentifiers") - - -def _backfill_scenario_identifiers() -> None: - """Backfill scenario rows and result foreign keys from the retained JSON column.""" - bind = op.get_bind() - result_rows = bind.execute( - sa.text('SELECT id, scenario_identifier FROM "ScenarioResultEntries" WHERE scenario_identifier IS NOT NULL') - ).fetchall() - update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') - inserter = IdentifierGraphInserter(bind=bind) - skipped = 0 - for result_id, raw_scenario in result_rows: - identifier = load_identifier(raw_scenario) - if identifier is None: - skipped += 1 - continue - try: - identifier_hash = inserter.insert_scenario(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) - except Exception: - skipped += 1 - logger.warning( - f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", - exc_info=True, - ) - - if skipped: - logger.warning(f"ScenarioIdentifiers backfill skipped {skipped} scenario result row(s)") diff --git a/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py b/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py deleted file mode 100644 index b114aea949..0000000000 --- a/pyrit/memory/alembic/versions/c8e1f3a5b7d9_add_converter_identifiers_table.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Persist converter identifiers as content-addressed rows. - -Revision ID: c8e1f3a5b7d9 -Revises: b7d9f1a3c5e7 -Create Date: 2026-07-13 15:00:00.000000 -""" - -from __future__ import annotations - -import logging -import uuid -from collections.abc import Sequence # noqa: TC003 -from typing import TYPE_CHECKING, Any - -import sqlalchemy as sa -from alembic import op -from sqlalchemy.dialects.sqlite import CHAR -from sqlalchemy.types import TypeDecorator, Uuid - -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier_list, - run_best_effort_backfill, -) - -if TYPE_CHECKING: - from sqlalchemy.engine import Dialect - - -class _CustomUUID(TypeDecorator[uuid.UUID]): - """Frozen UUID type matching ``PromptMemoryEntries.id`` across dialects.""" - - impl = CHAR - cache_ok = True - - def load_dialect_impl(self, dialect: Dialect) -> Any: - if dialect.name == "sqlite": - return dialect.type_descriptor(CHAR(36)) - return dialect.type_descriptor(Uuid()) - - def process_bind_param(self, value: Any, dialect: Any) -> str | None: - return str(value) if value is not None else None - - def process_result_value(self, value: Any, dialect: Any) -> uuid.UUID | None: - if value is None: - return None - if isinstance(value, uuid.UUID): - return value - return uuid.UUID(value) - - -revision: str = "c8e1f3a5b7d9" -down_revision: str | None = "b7d9f1a3c5e7" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -logger = logging.getLogger(__name__) - - -def upgrade() -> None: - """Apply this schema upgrade.""" - op.create_table( - "ConverterIdentifiers", - sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=True), - sa.Column("class_module", sa.String(), nullable=True), - sa.Column("identifier_json", sa.JSON(), nullable=True), - sa.Column("supported_input_types", sa.JSON(), nullable=True), - sa.Column("supported_output_types", sa.JSON(), nullable=True), - sa.Column("converter_target_hash", sa.String(64), nullable=True), - sa.Column("sub_converter_hash", sa.String(64), nullable=True), - sa.Column("pyrit_version", sa.String(), nullable=True), - sa.ForeignKeyConstraint( - ["converter_target_hash"], - ["TargetIdentifiers.hash"], - name="fk_converter_identifiers_converter_target_hash", - ), - sa.ForeignKeyConstraint( - ["sub_converter_hash"], - ["ConverterIdentifiers.hash"], - name="fk_converter_identifiers_sub_converter_hash", - ), - ) - op.create_table( - "PromptConverterIdentifiers", - sa.Column("prompt_memory_entry_id", _CustomUUID(), nullable=False), - sa.Column("position", sa.Integer(), nullable=False), - sa.Column("converter_identifier_hash", sa.String(64), nullable=False), - sa.ForeignKeyConstraint( - ["prompt_memory_entry_id"], - ["PromptMemoryEntries.id"], - name="fk_prompt_converter_identifiers_prompt_memory_entry_id", - ), - sa.ForeignKeyConstraint( - ["converter_identifier_hash"], - ["ConverterIdentifiers.hash"], - name="fk_prompt_converter_identifiers_converter_identifier_hash", - ), - sa.PrimaryKeyConstraint("prompt_memory_entry_id", "position"), - ) - bind = op.get_bind() - run_best_effort_backfill(bind=bind, name="ConverterIdentifiers", backfill=_backfill_converter_identifiers) - - -def downgrade() -> None: - """Revert this schema upgrade.""" - op.drop_table("PromptConverterIdentifiers") - op.drop_table("ConverterIdentifiers") - - -def _backfill_converter_identifiers() -> None: - """Materialize converter graphs and prompt associations from retained JSON.""" - bind = op.get_bind() - prompt_rows = bind.execute( - sa.text( - 'SELECT id, converter_identifiers, pyrit_version FROM "PromptMemoryEntries" ' - "WHERE converter_identifiers IS NOT NULL" - ) - ).fetchall() - link_insert = sa.text( - 'INSERT INTO "PromptConverterIdentifiers" ' - "(prompt_memory_entry_id, position, converter_identifier_hash) " - "VALUES (:prompt_memory_entry_id, :position, :converter_identifier_hash)" - ) - - inserter = IdentifierGraphInserter(bind=bind) - skipped = 0 - for prompt_id, stored_identifiers, pyrit_version in prompt_rows: - try: - for position, identifier in enumerate(load_identifier_list(stored_identifiers)): - if identifier.get("pyrit_version") is None: - identifier = {**identifier, "pyrit_version": pyrit_version} - identifier_hash = inserter.insert_converter(identifier) - if identifier_hash: - bind.execute( - link_insert, - { - "prompt_memory_entry_id": prompt_id, - "position": position, - "converter_identifier_hash": identifier_hash, - }, - ) - except Exception: - skipped += 1 - logger.warning( - f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", - exc_info=True, - ) - if skipped: - logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") diff --git a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py b/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py deleted file mode 100644 index 7e7bde32df..0000000000 --- a/pyrit/memory/alembic/versions/d9f2a4b6c8e0_add_attack_identifiers_tables.py +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Persist attack identifier graphs as content-addressed rows. - -Revision ID: d9f2a4b6c8e0 -Revises: c8e1f3a5b7d9 -Create Date: 2026-07-13 17:00:00.000000 -""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence # noqa: TC003 -from typing import Any - -import sqlalchemy as sa -from alembic import op - -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier, - run_best_effort_backfill, -) - -revision: str = "d9f2a4b6c8e0" -down_revision: str | None = "c8e1f3a5b7d9" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -logger = logging.getLogger(__name__) - - -def upgrade() -> None: - """Apply this schema upgrade.""" - _create_identifier_tables() - with op.batch_alter_table("AttackResultEntries") as batch_op: - batch_op.add_column(sa.Column("atomic_attack_identifier_hash", sa.String(64), nullable=True)) - batch_op.create_foreign_key( - "fk_attack_result_entries_atomic_attack_identifier_hash", - "AtomicAttackIdentifiers", - ["atomic_attack_identifier_hash"], - ["hash"], - ) - bind = op.get_bind() - run_best_effort_backfill(bind=bind, name="AttackIdentifiers", backfill=_backfill_attack_identifiers) - - -def downgrade() -> None: - """Revert this schema upgrade.""" - with op.batch_alter_table("AttackResultEntries") as batch_op: - batch_op.drop_column("atomic_attack_identifier_hash") - op.drop_table("AtomicAttackSeedIdentifiers") - op.drop_table("AtomicAttackIdentifiers") - op.drop_table("AttackTechniqueSeedIdentifiers") - op.drop_table("AttackTechniqueIdentifiers") - op.drop_table("AttackResponseConverterIdentifiers") - op.drop_table("AttackRequestConverterIdentifiers") - op.drop_table("AttackIdentifiers") - op.drop_table("SeedIdentifiers") - - -def _create_identifier_tables() -> None: - op.create_table( - "SeedIdentifiers", - *_common_columns(), - sa.Column("value", sa.Unicode(), nullable=True), - sa.Column("value_sha256", sa.String(), nullable=True), - sa.Column("data_type", sa.String(), nullable=True), - sa.Column("dataset_name", sa.String(), nullable=True), - sa.Column("is_general_technique", sa.Boolean(), nullable=True), - ) - op.create_table( - "AttackIdentifiers", - *_common_columns(), - sa.Column("adversarial_system_prompt", sa.Unicode(), nullable=True), - sa.Column("adversarial_seed_prompt", sa.Unicode(), nullable=True), - sa.Column("objective_target_hash", sa.String(64), nullable=True), - sa.Column("adversarial_chat_hash", sa.String(64), nullable=True), - sa.Column("objective_scorer_hash", sa.String(64), nullable=True), - sa.ForeignKeyConstraint(["objective_target_hash"], ["TargetIdentifiers.hash"]), - sa.ForeignKeyConstraint(["adversarial_chat_hash"], ["TargetIdentifiers.hash"]), - sa.ForeignKeyConstraint(["objective_scorer_hash"], ["ScorerIdentifiers.hash"]), - ) - _create_ordered_edge_table( - table_name="AttackRequestConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_table="AttackIdentifiers", - child_column="converter_identifier_hash", - child_table="ConverterIdentifiers", - ) - _create_ordered_edge_table( - table_name="AttackResponseConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_table="AttackIdentifiers", - child_column="converter_identifier_hash", - child_table="ConverterIdentifiers", - ) - op.create_table( - "AttackTechniqueIdentifiers", - *_common_columns(), - sa.Column("attack_identifier_hash", sa.String(64), nullable=True), - sa.ForeignKeyConstraint(["attack_identifier_hash"], ["AttackIdentifiers.hash"]), - ) - _create_ordered_edge_table( - table_name="AttackTechniqueSeedIdentifiers", - parent_column="attack_technique_identifier_hash", - parent_table="AttackTechniqueIdentifiers", - child_column="seed_identifier_hash", - child_table="SeedIdentifiers", - ) - op.create_table( - "AtomicAttackIdentifiers", - *_common_columns(), - sa.Column("attack_technique_identifier_hash", sa.String(64), nullable=True), - sa.ForeignKeyConstraint( - ["attack_technique_identifier_hash"], - ["AttackTechniqueIdentifiers.hash"], - ), - ) - _create_ordered_edge_table( - table_name="AtomicAttackSeedIdentifiers", - parent_column="atomic_attack_identifier_hash", - parent_table="AtomicAttackIdentifiers", - child_column="seed_identifier_hash", - child_table="SeedIdentifiers", - ) - - -def _common_columns() -> tuple[sa.Column[Any], ...]: - return ( - sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=True), - sa.Column("class_module", sa.String(), nullable=True), - sa.Column("identifier_json", sa.JSON(), nullable=True), - sa.Column("pyrit_version", sa.String(), nullable=True), - ) - - -def _create_ordered_edge_table( - *, - table_name: str, - parent_column: str, - parent_table: str, - child_column: str, - child_table: str, -) -> None: - op.create_table( - table_name, - sa.Column(parent_column, sa.String(64), nullable=False), - sa.Column("position", sa.Integer(), nullable=False), - sa.Column(child_column, sa.String(64), nullable=False), - sa.ForeignKeyConstraint([parent_column], [f"{parent_table}.hash"]), - sa.ForeignKeyConstraint([child_column], [f"{child_table}.hash"]), - sa.PrimaryKeyConstraint(parent_column, "position"), - ) - - -def _backfill_attack_identifiers() -> None: - """Backfill attack identifier graphs and result links from retained JSON.""" - bind = op.get_bind() - result_rows = bind.execute( - sa.text( - 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" WHERE atomic_attack_identifier IS NOT NULL' - ) - ).fetchall() - inserter = IdentifierGraphInserter(bind=bind) - update_stmt = sa.text('UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id') - - skipped = 0 - for result_id, raw_identifier in result_rows: - identifier = load_identifier(raw_identifier) - if identifier is None: - skipped += 1 - continue - try: - identifier_hash = inserter.insert_atomic_attack(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) - except Exception: - skipped += 1 - logger.warning( - f"Attack identifier backfill could not reconstruct result {result_id}", - exc_info=True, - ) - if skipped: - logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py new file mode 100644 index 0000000000..a70adc51f0 --- /dev/null +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py @@ -0,0 +1,519 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist component identifiers as content-addressed rows. + +Creates the normalized identifier tables, their graph edges, and nullable links +from existing domain tables. Retained identifier JSON is backfilled on a +best-effort basis and remains available when a legacy value cannot be linked. + +Revision ID: e5f7a9c1b3d2 +Revises: d4e6f8a0b2c4 +Create Date: 2026-07-10 12:00:00.000000 +""" + +from __future__ import annotations + +import logging +import uuid +from collections.abc import Sequence # noqa: TC003 +from typing import TYPE_CHECKING, Any + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.sqlite import CHAR +from sqlalchemy.types import TypeDecorator, Uuid + +from pyrit.memory.alembic.identifier_backfill import ( + IdentifierGraphInserter, + load_identifier, + load_identifier_list, + run_best_effort_backfill, +) + +if TYPE_CHECKING: + from sqlalchemy.engine import Dialect + +# revision identifiers, used by Alembic. +revision: str = "e5f7a9c1b3d2" +down_revision: str | None = "d4e6f8a0b2c4" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +class _CustomUUID(TypeDecorator[uuid.UUID]): + """Frozen UUID type matching ``PromptMemoryEntries.id`` across dialects.""" + + impl = CHAR + cache_ok = True + + def load_dialect_impl(self, dialect: Dialect) -> Any: + if dialect.name == "sqlite": + return dialect.type_descriptor(CHAR(36)) + return dialect.type_descriptor(Uuid()) + + def process_bind_param(self, value: Any, dialect: Any) -> str | None: + return str(value) if value is not None else None + + def process_result_value(self, value: Any, dialect: Any) -> uuid.UUID | None: + if value is None: + return None + if isinstance(value, uuid.UUID): + return value + return uuid.UUID(value) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "TargetIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), + sa.Column("endpoint", sa.String(), nullable=True), + sa.Column("model_name", sa.String(), nullable=True), + sa.Column("underlying_model_name", sa.String(), nullable=True), + sa.Column("temperature", sa.Float(), nullable=True), + sa.Column("top_p", sa.Float(), nullable=True), + sa.Column("max_requests_per_minute", sa.Integer(), nullable=True), + sa.Column("supported_auth_modes", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + # Self-referential pivot mapping a multi-target to its inner target identifiers. + # Both endpoints are content hashes into TargetIdentifiers; ``position`` preserves + # the parent's ``targets`` list order. Named FK constraints for SQL Server / batch + # portability. + op.create_table( + "TargetIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_child_hash" + ), + ) + + op.create_table( + "ScorerIdentifiers", + *_common_columns(), + sa.Column("scorer_type", sa.String(), nullable=True), + sa.Column("score_aggregator", sa.String(), nullable=True), + sa.Column("prompt_target_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["prompt_target_hash"], ["TargetIdentifiers.hash"], name="fk_scorer_identifiers_prompt_target_hash" + ), + ) + op.create_table( + "ScorerIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_child_hash" + ), + ) + op.create_table( + "ScenarioIdentifiers", + *_common_columns(), + sa.Column("version", sa.Integer(), nullable=True), + sa.Column("techniques", sa.JSON(), nullable=True), + sa.Column("datasets", sa.JSON(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["objective_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_scenario_identifiers_objective_target_hash", + ), + sa.ForeignKeyConstraint( + ["objective_scorer_hash"], + ["ScorerIdentifiers.hash"], + name="fk_scenario_identifiers_objective_scorer_hash", + ), + ) + op.create_table( + "ConverterIdentifiers", + *_common_columns(), + sa.Column("supported_input_types", sa.JSON(), nullable=True), + sa.Column("supported_output_types", sa.JSON(), nullable=True), + sa.Column("converter_target_hash", sa.String(64), nullable=True), + sa.Column("sub_converter_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["converter_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_converter_identifiers_converter_target_hash", + ), + sa.ForeignKeyConstraint( + ["sub_converter_hash"], + ["ConverterIdentifiers.hash"], + name="fk_converter_identifiers_sub_converter_hash", + ), + ) + op.create_table( + "PromptConverterIdentifiers", + sa.Column("prompt_memory_entry_id", _CustomUUID(), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("converter_identifier_hash", sa.String(64), nullable=False), + sa.ForeignKeyConstraint( + ["prompt_memory_entry_id"], + ["PromptMemoryEntries.id"], + name="fk_prompt_converter_identifiers_prompt_memory_entry_id", + ), + sa.ForeignKeyConstraint( + ["converter_identifier_hash"], + ["ConverterIdentifiers.hash"], + name="fk_prompt_converter_identifiers_converter_identifier_hash", + ), + sa.PrimaryKeyConstraint("prompt_memory_entry_id", "position"), + ) + _create_attack_identifier_tables() + + # Batch op for SQLite portability (no ALTER TABLE ADD FOREIGN KEY on SQLite). + # The FK constraint must be named explicitly: Alembic batch mode rejects an + # unnamed constraint. + with op.batch_alter_table("Conversations") as batch_op: + batch_op.add_column(sa.Column("target_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_conversations_target_identifier_hash", + "TargetIdentifiers", + ["target_identifier_hash"], + ["hash"], + ) + + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.add_column(sa.Column("scorer_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_score_entries_scorer_identifier_hash", + "ScorerIdentifiers", + ["scorer_identifier_hash"], + ["hash"], + ) + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.add_column(sa.Column("scenario_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_scenario_result_entries_scenario_identifier_hash", + "ScenarioIdentifiers", + ["scenario_identifier_hash"], + ["hash"], + ) + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.add_column(sa.Column("atomic_attack_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_attack_result_entries_atomic_attack_identifier_hash", + "AtomicAttackIdentifiers", + ["atomic_attack_identifier_hash"], + ["hash"], + ) + + bind = op.get_bind() + for name, backfill in ( + ("TargetIdentifiers", _backfill_target_identifiers), + ("ScorerIdentifiers", _backfill_scorer_identifiers), + ("ScenarioIdentifiers", _backfill_scenario_identifiers), + ("ConverterIdentifiers", _backfill_converter_identifiers), + ("AttackIdentifiers", _backfill_attack_identifiers), + ): + run_best_effort_backfill(bind=bind, name=name, backfill=backfill) + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.drop_column("atomic_attack_identifier_hash") + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.drop_column("scenario_identifier_hash") + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.drop_column("scorer_identifier_hash") + with op.batch_alter_table("Conversations") as batch_op: + batch_op.drop_column("target_identifier_hash") + + op.drop_table("AtomicAttackSeedIdentifiers") + op.drop_table("AtomicAttackIdentifiers") + op.drop_table("AttackTechniqueSeedIdentifiers") + op.drop_table("AttackTechniqueIdentifiers") + op.drop_table("AttackResponseConverterIdentifiers") + op.drop_table("AttackRequestConverterIdentifiers") + op.drop_table("AttackIdentifiers") + op.drop_table("SeedIdentifiers") + op.drop_table("PromptConverterIdentifiers") + op.drop_table("ConverterIdentifiers") + op.drop_table("ScenarioIdentifiers") + op.drop_table("ScorerIdentifierChildren") + op.drop_table("ScorerIdentifiers") + op.drop_table("TargetIdentifierChildren") + op.drop_table("TargetIdentifiers") + + +def _common_columns() -> tuple[sa.Column[Any], ...]: + return ( + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + +def _create_ordered_edge_table( + *, + table_name: str, + parent_column: str, + parent_table: str, + child_column: str, + child_table: str, +) -> None: + op.create_table( + table_name, + sa.Column(parent_column, sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column(child_column, sa.String(64), nullable=False), + sa.ForeignKeyConstraint([parent_column], [f"{parent_table}.hash"]), + sa.ForeignKeyConstraint([child_column], [f"{child_table}.hash"]), + sa.PrimaryKeyConstraint(parent_column, "position"), + ) + + +def _create_attack_identifier_tables() -> None: + op.create_table( + "SeedIdentifiers", + *_common_columns(), + sa.Column("value", sa.Unicode(), nullable=True), + sa.Column("value_sha256", sa.String(), nullable=True), + sa.Column("data_type", sa.String(), nullable=True), + sa.Column("dataset_name", sa.String(), nullable=True), + sa.Column("is_general_technique", sa.Boolean(), nullable=True), + ) + op.create_table( + "AttackIdentifiers", + *_common_columns(), + sa.Column("adversarial_system_prompt", sa.Unicode(), nullable=True), + sa.Column("adversarial_seed_prompt", sa.Unicode(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("adversarial_chat_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["objective_target_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["adversarial_chat_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["objective_scorer_hash"], ["ScorerIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + _create_ordered_edge_table( + table_name="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + op.create_table( + "AttackTechniqueIdentifiers", + *_common_columns(), + sa.Column("attack_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["attack_identifier_hash"], ["AttackIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_table="AttackTechniqueIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + op.create_table( + "AtomicAttackIdentifiers", + *_common_columns(), + sa.Column("attack_technique_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["attack_technique_identifier_hash"], ["AttackTechniqueIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_table="AtomicAttackIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + + +def _backfill_target_identifiers() -> None: + """ + Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set + ``Conversations.target_identifier_hash``. + + For every ``Conversations`` row with a non-null ``target_identifier`` JSON, + reconstruct the ``TargetIdentifier`` (recomputing its content hash), insert the + deduped ``TargetIdentifiers`` row if absent -- recursing into any inner + ``targets`` first so the child edge foreign keys resolve -- record the + ``parent_hash -> child_hash`` edges, and point the conversation's + ``target_identifier_hash`` at the top-level row. Idempotent: hashes already present + are not re-inserted. Rows whose stored target cannot be reconstructed are logged and + skipped rather than aborting the upgrade. + """ + bind = op.get_bind() + rows = bind.execute( + sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') + ).fetchall() + + update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') + inserter = IdentifierGraphInserter(bind=bind) + linked = 0 + skipped = 0 + for conversation_id, raw_target in rows: + identifier = load_identifier(raw_target) + if identifier is None: + skipped += 1 + continue + try: + identifier_hash = inserter.insert_target(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "cid": conversation_id}) + linked += 1 + except Exception: + skipped += 1 + logger.warning(f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", exc_info=True) + + if linked or skipped: + logger.info(f"TargetIdentifiers backfill linked {linked} conversation(s); skipped {skipped}.") + + +def _backfill_scorer_identifiers() -> None: + """Backfill scorer rows and score foreign keys from retained JSON.""" + bind = op.get_bind() + score_rows = bind.execute( + sa.text('SELECT id, scorer_class_identifier FROM "ScoreEntries" WHERE scorer_class_identifier IS NOT NULL') + ).fetchall() + score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for score_id, raw_scorer in score_rows: + identifier = load_identifier(raw_scorer) + if identifier is None: + skipped += 1 + continue + try: + identifier_hash = inserter.insert_scorer(identifier) + if identifier_hash: + bind.execute(score_update, {"hash": identifier_hash, "id": score_id}) + except Exception: + skipped += 1 + logger.warning( + f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", + exc_info=True, + ) + if skipped: + logger.warning(f"ScorerIdentifiers backfill skipped {skipped} score row(s)") + + +def _backfill_scenario_identifiers() -> None: + """Backfill scenario rows and result foreign keys from retained JSON.""" + bind = op.get_bind() + result_rows = bind.execute( + sa.text('SELECT id, scenario_identifier FROM "ScenarioResultEntries" WHERE scenario_identifier IS NOT NULL') + ).fetchall() + update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for result_id, raw_scenario in result_rows: + identifier = load_identifier(raw_scenario) + if identifier is None: + skipped += 1 + continue + try: + identifier_hash = inserter.insert_scenario(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) + except Exception: + skipped += 1 + logger.warning( + f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", + exc_info=True, + ) + if skipped: + logger.warning(f"ScenarioIdentifiers backfill skipped {skipped} scenario result row(s)") + + +def _backfill_converter_identifiers() -> None: + """Materialize converter graphs and prompt associations from retained JSON.""" + bind = op.get_bind() + prompt_rows = bind.execute( + sa.text( + 'SELECT id, converter_identifiers, pyrit_version FROM "PromptMemoryEntries" ' + "WHERE converter_identifiers IS NOT NULL" + ) + ).fetchall() + link_insert = sa.text( + 'INSERT INTO "PromptConverterIdentifiers" ' + "(prompt_memory_entry_id, position, converter_identifier_hash) " + "VALUES (:prompt_memory_entry_id, :position, :converter_identifier_hash)" + ) + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for prompt_id, stored_identifiers, pyrit_version in prompt_rows: + try: + for position, identifier in enumerate(load_identifier_list(stored_identifiers)): + if identifier.get("pyrit_version") is None: + identifier = {**identifier, "pyrit_version": pyrit_version} + identifier_hash = inserter.insert_converter(identifier) + if identifier_hash: + bind.execute( + link_insert, + { + "prompt_memory_entry_id": prompt_id, + "position": position, + "converter_identifier_hash": identifier_hash, + }, + ) + except Exception: + skipped += 1 + logger.warning( + f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", + exc_info=True, + ) + if skipped: + logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") + + +def _backfill_attack_identifiers() -> None: + """Backfill attack identifier graphs and result links from retained JSON.""" + bind = op.get_bind() + result_rows = bind.execute( + sa.text( + 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" WHERE atomic_attack_identifier IS NOT NULL' + ) + ).fetchall() + inserter = IdentifierGraphInserter(bind=bind) + update_stmt = sa.text('UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id') + skipped = 0 + for result_id, raw_identifier in result_rows: + identifier = load_identifier(raw_identifier) + if identifier is None: + skipped += 1 + continue + try: + identifier_hash = inserter.insert_atomic_attack(identifier) + if identifier_hash: + bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) + except Exception: + skipped += 1 + logger.warning( + f"Attack identifier backfill could not reconstruct result {result_id}", + exc_info=True, + ) + if skipped: + logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py deleted file mode 100644 index 4ed290f2ed..0000000000 --- a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py +++ /dev/null @@ -1,149 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Introduce the TargetIdentifiers table and reference it from Conversations. - -Phase 1 (dual-write) of storing component identifiers as first-class, -content-addressed rows. Creates ``TargetIdentifiers`` (one row per distinct -target identifier, keyed by its content ``hash``, with promoted scalar query -columns) and ``TargetIdentifierChildren`` (a self-referential pivot mapping a -multi-target to its inner target identifiers), adds a nullable -``target_identifier_hash`` foreign key to ``Conversations``, and backfills all -three from the existing ``Conversations.target_identifier`` JSON column. The JSON -column is retained (reads still come from it), so this migration is purely -additive. - -Revision ID: e5f7a9c1b3d2 -Revises: d4e6f8a0b2c4 -Create Date: 2026-07-10 12:00:00.000000 -""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence # noqa: TC003 - -import sqlalchemy as sa -from alembic import op - -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier, - run_best_effort_backfill, -) - -# revision identifiers, used by Alembic. -revision: str = "e5f7a9c1b3d2" -down_revision: str | None = "d4e6f8a0b2c4" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - - -logger = logging.getLogger(__name__) - - -def upgrade() -> None: - """Apply this schema upgrade.""" - op.create_table( - "TargetIdentifiers", - sa.Column("hash", sa.String(64), primary_key=True, nullable=False), - sa.Column("class_name", sa.String(), nullable=True), - sa.Column("class_module", sa.String(), nullable=True), - sa.Column("identifier_json", sa.JSON(), nullable=True), - sa.Column("endpoint", sa.String(), nullable=True), - sa.Column("model_name", sa.String(), nullable=True), - sa.Column("underlying_model_name", sa.String(), nullable=True), - sa.Column("temperature", sa.Float(), nullable=True), - sa.Column("top_p", sa.Float(), nullable=True), - sa.Column("max_requests_per_minute", sa.Integer(), nullable=True), - sa.Column("supported_auth_modes", sa.JSON(), nullable=True), - sa.Column("pyrit_version", sa.String(), nullable=True), - ) - - # Self-referential pivot mapping a multi-target to its inner target identifiers. - # Both endpoints are content hashes into TargetIdentifiers; ``position`` preserves - # the parent's ``targets`` list order. Named FK constraints for SQL Server / batch - # portability. - op.create_table( - "TargetIdentifierChildren", - sa.Column("parent_hash", sa.String(64), nullable=False), - sa.Column("position", sa.Integer(), nullable=False), - sa.Column("child_hash", sa.String(64), nullable=False), - sa.PrimaryKeyConstraint("parent_hash", "position"), - sa.ForeignKeyConstraint( - ["parent_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_parent_hash" - ), - sa.ForeignKeyConstraint( - ["child_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_child_hash" - ), - ) - - # Batch op for SQLite portability (no ALTER TABLE ADD FOREIGN KEY on SQLite). - # The FK constraint must be named explicitly: Alembic batch mode rejects an - # unnamed constraint. - with op.batch_alter_table("Conversations") as batch_op: - batch_op.add_column(sa.Column("target_identifier_hash", sa.String(64), nullable=True)) - batch_op.create_foreign_key( - "fk_conversations_target_identifier_hash", - "TargetIdentifiers", - ["target_identifier_hash"], - ["hash"], - ) - - bind = op.get_bind() - run_best_effort_backfill( - bind=bind, - name="TargetIdentifiers", - backfill=_backfill_target_identifiers, - ) - - -def downgrade() -> None: - """Revert this schema upgrade.""" - with op.batch_alter_table("Conversations") as batch_op: - batch_op.drop_column("target_identifier_hash") - # Drop the child edge table before its referenced parent table. - op.drop_table("TargetIdentifierChildren") - op.drop_table("TargetIdentifiers") - - -def _backfill_target_identifiers() -> None: - """ - Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set - ``Conversations.target_identifier_hash``. - - For every ``Conversations`` row with a non-null ``target_identifier`` JSON, - reconstruct the ``TargetIdentifier`` (recomputing its content hash), insert the - deduped ``TargetIdentifiers`` row if absent -- recursing into any inner - ``targets`` first so the child edge foreign keys resolve -- record the - ``parent_hash -> child_hash`` edges, and point the conversation's - ``target_identifier_hash`` at the top-level row. Idempotent: hashes already present - are not re-inserted. Rows whose stored target cannot be reconstructed are logged and - skipped rather than aborting the upgrade. - """ - bind = op.get_bind() - rows = bind.execute( - sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') - ).fetchall() - - update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') - inserter = IdentifierGraphInserter(bind=bind) - linked = 0 - skipped = 0 - for conversation_id, raw_target in rows: - identifier = load_identifier(raw_target) - if identifier is None: - skipped += 1 - continue - try: - identifier_hash = inserter.insert_target(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "cid": conversation_id}) - linked += 1 - except Exception: - skipped += 1 - logger.warning(f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", exc_info=True) - - if linked or skipped: - logger.info(f"TargetIdentifiers backfill linked {linked} conversation(s); skipped {skipped}.") diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 715fb96389..3b89efd51c 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -810,7 +810,7 @@ def test_conversations_migration_downgrade_restores_columns(): # ============================================================================= -# Backfill tests for scorer identifier persistence (a6c8e0f2b4d6) +# Backfill tests for identifier persistence (e5f7a9c1b3d2) # ============================================================================= @@ -826,17 +826,6 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( text('INSERT INTO "Conversations" (conversation_id, target_identifier) VALUES (:id, :value)'), {"id": "malformed-conversation", "value": "not-json"}, ) - command.upgrade(config, "e5f7a9c1b3d2") - assert ( - connection.execute( - text( - 'SELECT target_identifier_hash FROM "Conversations" ' - "WHERE conversation_id = 'malformed-conversation'" - ) - ).scalar_one() - is None - ) - score_id = str(uuid.uuid4()) connection.execute( text( @@ -846,15 +835,6 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( ), {"id": score_id}, ) - command.upgrade(config, "a6c8e0f2b4d6") - assert ( - connection.execute( - text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), - {"id": score_id}, - ).scalar_one() - is None - ) - result_id = str(uuid.uuid4()) connection.execute( text( @@ -866,15 +846,6 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( ), {"id": result_id}, ) - command.upgrade(config, "b7d9f1a3c5e7") - assert ( - connection.execute( - text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), - {"id": result_id}, - ).scalar_one() - is None - ) - prompt_id = str(uuid.uuid4()) connection.execute( text( @@ -886,9 +857,6 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( ), {"id": prompt_id}, ) - command.upgrade(config, "c8e1f3a5b7d9") - assert connection.execute(text('SELECT COUNT(*) FROM "PromptConverterIdentifiers"')).scalar_one() == 0 - attack_result_id = str(uuid.uuid4()) connection.execute( text( @@ -899,7 +867,33 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( ), {"id": attack_result_id}, ) - command.upgrade(config, "d9f2a4b6c8e0") + + command.upgrade(config, "e5f7a9c1b3d2") + + assert ( + connection.execute( + text( + 'SELECT target_identifier_hash FROM "Conversations" ' + "WHERE conversation_id = 'malformed-conversation'" + ) + ).scalar_one() + is None + ) + assert ( + connection.execute( + text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), + {"id": score_id}, + ).scalar_one() + is None + ) + assert ( + connection.execute( + text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() + is None + ) + assert connection.execute(text('SELECT COUNT(*) FROM "PromptConverterIdentifiers"')).scalar_one() == 0 assert ( connection.execute( text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), @@ -930,13 +924,7 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( def test_identifier_migrations_do_not_import_domain_models(): """Frozen identifier migrations operate on retained JSON rather than current domain models.""" versions_dir = Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic" / "versions" - revision_names = ( - "e5f7a9c1b3d2_add_target_identifiers_table.py", - "a6c8e0f2b4d6_add_scorer_identifiers_table.py", - "b7d9f1a3c5e7_add_scenario_identifiers_table.py", - "c8e1f3a5b7d9_add_converter_identifiers_table.py", - "d9f2a4b6c8e0_add_attack_identifiers_tables.py", - ) + revision_names = ("e5f7a9c1b3d2_add_identifiers_tables.py",) for revision_name in revision_names: source = (versions_dir / revision_name).read_text(encoding="utf-8") @@ -971,7 +959,7 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): try: with engine.begin() as connection: config = _config_for(connection) - command.upgrade(config, "e5f7a9c1b3d2") + command.upgrade(config, "d4e6f8a0b2c4") connection.execute( text( 'INSERT INTO "ScoreEntries" ' @@ -981,7 +969,7 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): {"id": score_id, "identifier": json.dumps(composite.model_dump())}, ) - command.upgrade(config, "a6c8e0f2b4d6") + command.upgrade(config, "e5f7a9c1b3d2") score_hash = connection.execute( text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), @@ -1008,7 +996,7 @@ def test_scorer_identifier_migration_backfills_graph_and_score_link(): # ============================================================================= -# Backfill tests for converter identifier persistence (c8e1f3a5b7d9) +# Backfill tests for converter identifier persistence # ============================================================================= @@ -1042,7 +1030,7 @@ def test_converter_identifier_migration_backfills_graph_and_prompt_links(): try: with engine.begin() as connection: config = _config_for(connection) - command.upgrade(config, "b7d9f1a3c5e7") + command.upgrade(config, "d4e6f8a0b2c4") connection.execute( text( 'INSERT INTO "PromptMemoryEntries" ' @@ -1058,7 +1046,7 @@ def test_converter_identifier_migration_backfills_graph_and_prompt_links(): }, ) - command.upgrade(config, "c8e1f3a5b7d9") + command.upgrade(config, "e5f7a9c1b3d2") converter_rows = connection.execute( text('SELECT hash, converter_target_hash, sub_converter_hash FROM "ConverterIdentifiers"') @@ -1082,7 +1070,7 @@ def test_converter_identifier_migration_backfills_graph_and_prompt_links(): # ============================================================================= -# Backfill tests for scenario identifier persistence (b7d9f1a3c5e7) +# Backfill tests for scenario identifier persistence # ============================================================================= @@ -1122,7 +1110,7 @@ def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): try: with engine.begin() as connection: config = _config_for(connection) - command.upgrade(config, "a6c8e0f2b4d6") + command.upgrade(config, "d4e6f8a0b2c4") connection.execute( text( 'INSERT INTO "ScenarioResultEntries" ' @@ -1140,7 +1128,7 @@ def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): }, ) - command.upgrade(config, "b7d9f1a3c5e7") + command.upgrade(config, "e5f7a9c1b3d2") result_hash = connection.execute( text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), @@ -1171,7 +1159,7 @@ def test_scenario_identifier_migration_backfills_dependencies_and_result_link(): # ============================================================================= -# Backfill tests for attack identifier persistence (d9f2a4b6c8e0) +# Backfill tests for attack identifier persistence # ============================================================================= @@ -1233,7 +1221,7 @@ def test_attack_identifier_migration_backfills_graph_and_result_link(): try: with engine.begin() as connection: config = _config_for(connection) - command.upgrade(config, "c8e1f3a5b7d9") + command.upgrade(config, "d4e6f8a0b2c4") connection.execute( text( 'INSERT INTO "AttackResultEntries" ' @@ -1245,7 +1233,7 @@ def test_attack_identifier_migration_backfills_graph_and_result_link(): {"id": result_id, "identifier": json.dumps(atomic.model_dump())}, ) - command.upgrade(config, "d9f2a4b6c8e0") + command.upgrade(config, "e5f7a9c1b3d2") result_hash = connection.execute( text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), From 888ed0ad59bc2fe66ef8f85e708427ec86026ed3 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 11:37:38 -0700 Subject: [PATCH 15/24] helpers in migration script --- pyrit/memory/alembic/identifier_backfill.py | 358 ------------------ .../e5f7a9c1b3d2_add_identifiers_tables.py | 351 ++++++++++++++++- 2 files changed, 344 insertions(+), 365 deletions(-) delete mode 100644 pyrit/memory/alembic/identifier_backfill.py diff --git a/pyrit/memory/alembic/identifier_backfill.py b/pyrit/memory/alembic/identifier_backfill.py deleted file mode 100644 index 3a3b22c522..0000000000 --- a/pyrit/memory/alembic/identifier_backfill.py +++ /dev/null @@ -1,358 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Frozen, model-independent helpers for identifier migration backfills.""" - -from __future__ import annotations - -import json -import logging -from typing import TYPE_CHECKING, Any - -import sqlalchemy as sa - -if TYPE_CHECKING: - from collections.abc import Callable, Sequence - -logger = logging.getLogger(__name__) - - -def run_best_effort_backfill(*, bind: Any, name: str, backfill: Callable[[], None]) -> None: - """Run a data backfill in a savepoint without blocking the schema upgrade.""" - try: - with bind.begin_nested(): - backfill() - except Exception: - logger.warning(f"{name} backfill failed; leaving new identifier links nullable", exc_info=True) - - -def load_identifier(raw_identifier: Any) -> dict[str, Any] | None: - """ - Load a retained identifier JSON value without importing domain models. - - Returns: - dict[str, Any] | None: The identifier dictionary when it has a usable hash. - """ - try: - value = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier - except (TypeError, ValueError): - return None - if not isinstance(value, dict): - return None - identifier_hash = value.get("hash") - if not isinstance(identifier_hash, str) or len(identifier_hash) != 64: - return None - return value - - -def load_identifier_list(raw_identifiers: Any) -> list[dict[str, Any]]: - """ - Load the valid identifiers from a retained JSON list. - - Returns: - list[dict[str, Any]]: Identifier dictionaries carrying usable hashes. - """ - try: - values = json.loads(raw_identifiers) if isinstance(raw_identifiers, str) else raw_identifiers - except (TypeError, ValueError): - return [] - if not isinstance(values, list): - return [] - return [identifier for value in values if (identifier := load_identifier(value)) is not None] - - -class IdentifierGraphInserter: - """Best-effort inserter for the frozen flat identifier JSON shape.""" - - _TABLES = ( - "TargetIdentifiers", - "ScorerIdentifiers", - "ConverterIdentifiers", - "ScenarioIdentifiers", - "SeedIdentifiers", - "AttackIdentifiers", - "AttackTechniqueIdentifiers", - "AtomicAttackIdentifiers", - ) - - def __init__(self, *, bind: Any) -> None: - """Initialize the inserter from tables available at this migration revision.""" - self._bind = bind - table_names = set(sa.inspect(bind).get_table_names()) - self._hashes = { - table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) - for table in self._TABLES - if table in table_names - } - - def insert_target(self, identifier: dict[str, Any]) -> str | None: - """ - Insert a target graph. - - Returns: - str | None: The stored hash when successful. - """ - children = self._children(identifier, "targets") - child_hashes = [child_hash for child in children if (child_hash := self.insert_target(child))] - identifier_hash = self._insert_identifier( - table="TargetIdentifiers", - identifier=identifier, - promoted=( - "endpoint", - "model_name", - "underlying_model_name", - "temperature", - "top_p", - "max_requests_per_minute", - "supported_auth_modes", - ), - ) - if identifier_hash: - self._insert_edges( - table="TargetIdentifierChildren", - parent_column="parent_hash", - parent_hash=identifier_hash, - child_column="child_hash", - child_hashes=child_hashes, - ) - return identifier_hash - - def insert_scorer(self, identifier: dict[str, Any]) -> str | None: - """ - Insert a scorer graph. - - Returns: - str | None: The stored hash when successful. - """ - prompt_target = self._child(identifier, "prompt_target", aliases=("chat_target",)) - prompt_target_hash = self.insert_target(prompt_target) if prompt_target else None - sub_scorers = self._children(identifier, "sub_scorers", aliases=("scorers",)) - child_hashes = [child_hash for child in sub_scorers if (child_hash := self.insert_scorer(child))] - identifier_hash = self._insert_identifier( - table="ScorerIdentifiers", - identifier=identifier, - promoted=("scorer_type", "score_aggregator"), - extra={"prompt_target_hash": prompt_target_hash}, - ) - if identifier_hash: - self._insert_edges( - table="ScorerIdentifierChildren", - parent_column="parent_hash", - parent_hash=identifier_hash, - child_column="child_hash", - child_hashes=child_hashes, - ) - return identifier_hash - - def insert_converter(self, identifier: dict[str, Any]) -> str | None: - """ - Insert a converter graph. - - Returns: - str | None: The stored hash when successful. - """ - converter_target = self._child(identifier, "converter_target") - sub_converter = self._child(identifier, "sub_converter") - return self._insert_identifier( - table="ConverterIdentifiers", - identifier=identifier, - promoted=("supported_input_types", "supported_output_types"), - extra={ - "converter_target_hash": self.insert_target(converter_target) if converter_target else None, - "sub_converter_hash": self.insert_converter(sub_converter) if sub_converter else None, - }, - ) - - def insert_scenario(self, identifier: dict[str, Any]) -> str | None: - """ - Insert a scenario graph. - - Returns: - str | None: The stored hash when successful. - """ - objective_target = self._child(identifier, "objective_target") - objective_scorer = self._child(identifier, "objective_scorer") - return self._insert_identifier( - table="ScenarioIdentifiers", - identifier=identifier, - promoted=("version", "techniques", "datasets"), - extra={ - "objective_target_hash": self.insert_target(objective_target) if objective_target else None, - "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, - }, - ) - - def insert_atomic_attack(self, identifier: dict[str, Any]) -> str | None: - """ - Insert an atomic attack graph. - - Returns: - str | None: The stored hash when successful. - """ - attack_technique = self._child(identifier, "attack_technique") - seeds = self._children(identifier, "seed_identifiers") - seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] - identifier_hash = self._insert_identifier( - table="AtomicAttackIdentifiers", - identifier=identifier, - extra={ - "attack_technique_identifier_hash": ( - self._insert_attack_technique(attack_technique) if attack_technique else None - ) - }, - ) - if identifier_hash: - self._insert_edges( - table="AtomicAttackSeedIdentifiers", - parent_column="atomic_attack_identifier_hash", - parent_hash=identifier_hash, - child_column="seed_identifier_hash", - child_hashes=seed_hashes, - ) - return identifier_hash - - def _insert_attack_technique(self, identifier: dict[str, Any]) -> str | None: - attack = self._child(identifier, "attack") - seeds = self._children(identifier, "technique_seeds") - seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] - identifier_hash = self._insert_identifier( - table="AttackTechniqueIdentifiers", - identifier=identifier, - extra={"attack_identifier_hash": self._insert_attack(attack) if attack else None}, - ) - if identifier_hash: - self._insert_edges( - table="AttackTechniqueSeedIdentifiers", - parent_column="attack_technique_identifier_hash", - parent_hash=identifier_hash, - child_column="seed_identifier_hash", - child_hashes=seed_hashes, - ) - return identifier_hash - - def _insert_attack(self, identifier: dict[str, Any]) -> str | None: - objective_target = self._child(identifier, "objective_target") - adversarial_chat = self._child(identifier, "adversarial_chat") - objective_scorer = self._child(identifier, "objective_scorer") - request_hashes = [ - value for item in self._children(identifier, "request_converters") if (value := self.insert_converter(item)) - ] - response_hashes = [ - value - for item in self._children(identifier, "response_converters") - if (value := self.insert_converter(item)) - ] - identifier_hash = self._insert_identifier( - table="AttackIdentifiers", - identifier=identifier, - promoted=("adversarial_system_prompt", "adversarial_seed_prompt"), - extra={ - "objective_target_hash": self.insert_target(objective_target) if objective_target else None, - "adversarial_chat_hash": self.insert_target(adversarial_chat) if adversarial_chat else None, - "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, - }, - ) - if identifier_hash: - self._insert_edges( - table="AttackRequestConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_hash=identifier_hash, - child_column="converter_identifier_hash", - child_hashes=request_hashes, - ) - self._insert_edges( - table="AttackResponseConverterIdentifiers", - parent_column="attack_identifier_hash", - parent_hash=identifier_hash, - child_column="converter_identifier_hash", - child_hashes=response_hashes, - ) - return identifier_hash - - def _insert_seed(self, identifier: dict[str, Any]) -> str | None: - return self._insert_identifier( - table="SeedIdentifiers", - identifier=identifier, - promoted=("value", "value_sha256", "data_type", "dataset_name", "is_general_technique"), - ) - - def _insert_identifier( - self, - *, - table: str, - identifier: dict[str, Any], - promoted: Sequence[str] = (), - extra: dict[str, Any] | None = None, - ) -> str | None: - identifier_hash = identifier.get("hash") - if not isinstance(identifier_hash, str) or len(identifier_hash) != 64 or table not in self._hashes: - return None - if identifier_hash in self._hashes[table]: - return identifier_hash - values: dict[str, Any] = { - "hash": identifier_hash, - "class_name": identifier.get("class_name"), - "class_module": identifier.get("class_module"), - "identifier_json": json.dumps(identifier, sort_keys=True), - "pyrit_version": identifier.get("pyrit_version"), - } - values.update({name: self._json_value(identifier.get(name)) for name in promoted}) - values.update(extra or {}) - columns = list(values) - statement = sa.text( - f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(f":{column}" for column in columns)})' - ) - self._bind.execute(statement, values) - self._hashes[table].add(identifier_hash) - return identifier_hash - - def _insert_edges( - self, - *, - table: str, - parent_column: str, - parent_hash: str, - child_column: str, - child_hashes: Sequence[str], - ) -> None: - statement = sa.text( - f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' - f"VALUES (:parent_hash, :position, :child_hash)" - ) - for position, child_hash in enumerate(child_hashes): - self._bind.execute( - statement, - {"parent_hash": parent_hash, "position": position, "child_hash": child_hash}, - ) - - @staticmethod - def _child( - identifier: dict[str, Any], - name: str, - aliases: Sequence[str] = (), - ) -> dict[str, Any] | None: - children = identifier.get("children") - children = children if isinstance(children, dict) else {} - for key in (name, *aliases): - value = identifier.get(key, children.get(key)) - if isinstance(value, dict): - return load_identifier(value) - return None - - @staticmethod - def _children( - identifier: dict[str, Any], - name: str, - aliases: Sequence[str] = (), - ) -> list[dict[str, Any]]: - children = identifier.get("children") - children = children if isinstance(children, dict) else {} - for key in (name, *aliases): - value = identifier.get(key, children.get(key)) - if isinstance(value, list): - return [child for item in value if (child := load_identifier(item)) is not None] - return [] - - @staticmethod - def _json_value(value: Any) -> Any: - return json.dumps(value) if isinstance(value, (list, dict)) else value diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py index a70adc51f0..9d66a62a90 100644 --- a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py @@ -15,6 +15,7 @@ from __future__ import annotations +import json import logging import uuid from collections.abc import Sequence # noqa: TC003 @@ -25,14 +26,9 @@ from sqlalchemy.dialects.sqlite import CHAR from sqlalchemy.types import TypeDecorator, Uuid -from pyrit.memory.alembic.identifier_backfill import ( - IdentifierGraphInserter, - load_identifier, - load_identifier_list, - run_best_effort_backfill, -) - if TYPE_CHECKING: + from collections.abc import Callable + from sqlalchemy.engine import Dialect # revision identifiers, used by Alembic. @@ -67,6 +63,347 @@ def process_result_value(self, value: Any, dialect: Any) -> uuid.UUID | None: return uuid.UUID(value) +def run_best_effort_backfill(*, bind: Any, name: str, backfill: Callable[[], None]) -> None: + """Run a data backfill in a savepoint without blocking the schema upgrade.""" + try: + with bind.begin_nested(): + backfill() + except Exception: + logger.warning(f"{name} backfill failed; leaving new identifier links nullable", exc_info=True) + + +def load_identifier(raw_identifier: Any) -> dict[str, Any] | None: + """ + Load a retained identifier JSON value without importing domain models. + + Returns: + dict[str, Any] | None: The identifier dictionary when it has a usable hash. + """ + try: + value = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier + except (TypeError, ValueError): + return None + if not isinstance(value, dict): + return None + identifier_hash = value.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64: + return None + return value + + +def load_identifier_list(raw_identifiers: Any) -> list[dict[str, Any]]: + """ + Load the valid identifiers from a retained JSON list. + + Returns: + list[dict[str, Any]]: Identifier dictionaries carrying usable hashes. + """ + try: + values = json.loads(raw_identifiers) if isinstance(raw_identifiers, str) else raw_identifiers + except (TypeError, ValueError): + return [] + if not isinstance(values, list): + return [] + return [identifier for value in values if (identifier := load_identifier(value)) is not None] + + +class IdentifierGraphInserter: + """Best-effort inserter for the frozen flat identifier JSON shape.""" + + _TABLES = ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ConverterIdentifiers", + "ScenarioIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ) + + def __init__(self, *, bind: Any) -> None: + """Initialize the inserter from tables available at this migration revision.""" + self._bind = bind + table_names = set(sa.inspect(bind).get_table_names()) + self._hashes = { + table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) + for table in self._TABLES + if table in table_names + } + + def insert_target(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a target graph. + + Returns: + str | None: The stored hash when successful. + """ + children = self._children(identifier, "targets") + child_hashes = [child_hash for child in children if (child_hash := self.insert_target(child))] + identifier_hash = self._insert_identifier( + table="TargetIdentifiers", + identifier=identifier, + promoted=( + "endpoint", + "model_name", + "underlying_model_name", + "temperature", + "top_p", + "max_requests_per_minute", + "supported_auth_modes", + ), + ) + if identifier_hash: + self._insert_edges( + table="TargetIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_scorer(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scorer graph. + + Returns: + str | None: The stored hash when successful. + """ + prompt_target = self._child(identifier, "prompt_target", aliases=("chat_target",)) + prompt_target_hash = self.insert_target(prompt_target) if prompt_target else None + sub_scorers = self._children(identifier, "sub_scorers", aliases=("scorers",)) + child_hashes = [child_hash for child in sub_scorers if (child_hash := self.insert_scorer(child))] + identifier_hash = self._insert_identifier( + table="ScorerIdentifiers", + identifier=identifier, + promoted=("scorer_type", "score_aggregator"), + extra={"prompt_target_hash": prompt_target_hash}, + ) + if identifier_hash: + self._insert_edges( + table="ScorerIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_converter(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a converter graph. + + Returns: + str | None: The stored hash when successful. + """ + converter_target = self._child(identifier, "converter_target") + sub_converter = self._child(identifier, "sub_converter") + return self._insert_identifier( + table="ConverterIdentifiers", + identifier=identifier, + promoted=("supported_input_types", "supported_output_types"), + extra={ + "converter_target_hash": self.insert_target(converter_target) if converter_target else None, + "sub_converter_hash": self.insert_converter(sub_converter) if sub_converter else None, + }, + ) + + def insert_scenario(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scenario graph. + + Returns: + str | None: The stored hash when successful. + """ + objective_target = self._child(identifier, "objective_target") + objective_scorer = self._child(identifier, "objective_scorer") + return self._insert_identifier( + table="ScenarioIdentifiers", + identifier=identifier, + promoted=("version", "techniques", "datasets"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + + def insert_atomic_attack(self, identifier: dict[str, Any]) -> str | None: + """ + Insert an atomic attack graph. + + Returns: + str | None: The stored hash when successful. + """ + attack_technique = self._child(identifier, "attack_technique") + seeds = self._children(identifier, "seed_identifiers") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AtomicAttackIdentifiers", + identifier=identifier, + extra={ + "attack_technique_identifier_hash": ( + self._insert_attack_technique(attack_technique) if attack_technique else None + ) + }, + ) + if identifier_hash: + self._insert_edges( + table="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack_technique(self, identifier: dict[str, Any]) -> str | None: + attack = self._child(identifier, "attack") + seeds = self._children(identifier, "technique_seeds") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AttackTechniqueIdentifiers", + identifier=identifier, + extra={"attack_identifier_hash": self._insert_attack(attack) if attack else None}, + ) + if identifier_hash: + self._insert_edges( + table="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack(self, identifier: dict[str, Any]) -> str | None: + objective_target = self._child(identifier, "objective_target") + adversarial_chat = self._child(identifier, "adversarial_chat") + objective_scorer = self._child(identifier, "objective_scorer") + request_hashes = [ + value for item in self._children(identifier, "request_converters") if (value := self.insert_converter(item)) + ] + response_hashes = [ + value + for item in self._children(identifier, "response_converters") + if (value := self.insert_converter(item)) + ] + identifier_hash = self._insert_identifier( + table="AttackIdentifiers", + identifier=identifier, + promoted=("adversarial_system_prompt", "adversarial_seed_prompt"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "adversarial_chat_hash": self.insert_target(adversarial_chat) if adversarial_chat else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + if identifier_hash: + self._insert_edges( + table="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=request_hashes, + ) + self._insert_edges( + table="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=response_hashes, + ) + return identifier_hash + + def _insert_seed(self, identifier: dict[str, Any]) -> str | None: + return self._insert_identifier( + table="SeedIdentifiers", + identifier=identifier, + promoted=("value", "value_sha256", "data_type", "dataset_name", "is_general_technique"), + ) + + def _insert_identifier( + self, + *, + table: str, + identifier: dict[str, Any], + promoted: Sequence[str] = (), + extra: dict[str, Any] | None = None, + ) -> str | None: + identifier_hash = identifier.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64 or table not in self._hashes: + return None + if identifier_hash in self._hashes[table]: + return identifier_hash + values: dict[str, Any] = { + "hash": identifier_hash, + "class_name": identifier.get("class_name"), + "class_module": identifier.get("class_module"), + "identifier_json": json.dumps(identifier, sort_keys=True), + "pyrit_version": identifier.get("pyrit_version"), + } + values.update({name: self._json_value(identifier.get(name)) for name in promoted}) + values.update(extra or {}) + columns = list(values) + statement = sa.text( + f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(f":{column}" for column in columns)})' + ) + self._bind.execute(statement, values) + self._hashes[table].add(identifier_hash) + return identifier_hash + + def _insert_edges( + self, + *, + table: str, + parent_column: str, + parent_hash: str, + child_column: str, + child_hashes: Sequence[str], + ) -> None: + statement = sa.text( + f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' + f"VALUES (:parent_hash, :position, :child_hash)" + ) + for position, child_hash in enumerate(child_hashes): + self._bind.execute( + statement, + {"parent_hash": parent_hash, "position": position, "child_hash": child_hash}, + ) + + @staticmethod + def _child( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> dict[str, Any] | None: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, dict): + return load_identifier(value) + return None + + @staticmethod + def _children( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> list[dict[str, Any]]: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, list): + return [child for item in value if (child := load_identifier(item)) is not None] + return [] + + @staticmethod + def _json_value(value: Any) -> Any: + return json.dumps(value) if isinstance(value, (list, dict)) else value + + def upgrade() -> None: """Apply this schema upgrade.""" op.create_table( From 7d727faf5d98d90d2a68f706a1c4d6befa68190f Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 12:20:16 -0700 Subject: [PATCH 16/24] test --- tests/unit/memory/test_memory_models.py | 27 ++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 52107984d0..3c2c357346 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -10,7 +10,7 @@ import pytest from pydantic import ValidationError from sqlalchemy import create_engine, select -from sqlalchemy.orm import Session +from sqlalchemy.orm import MappedColumn, Session from pyrit.memory.memory_models import ( AtomicAttackIdentifierEntry, @@ -245,6 +245,31 @@ def test_identifier_entry_maps_promoted_children_by_cardinality( assert set(entry_type.CHILD_HASH_COLUMNS) == singular_children +@pytest.mark.parametrize( + ("identifier_type", "entry_type"), + [ + (TargetIdentifier, TargetIdentifierEntry), + (ConverterIdentifier, ConverterIdentifierEntry), + (ScorerIdentifier, ScorerIdentifierEntry), + (ScenarioIdentifier, ScenarioIdentifierEntry), + (SeedIdentifier, SeedIdentifierEntry), + (AttackIdentifier, AttackIdentifierEntry), + (AttackTechniqueIdentifier, AttackTechniqueIdentifierEntry), + (AtomicAttackIdentifier, AtomicAttackIdentifierEntry), + ], +) +def test_identifier_entry_maps_promoted_scalars_to_columns( + identifier_type: type[ComponentIdentifier], + entry_type: type[ComponentIdentifierEntry[Any]], +) -> None: + shared_columns = {name for name, value in vars(ComponentIdentifierEntry).items() if isinstance(value, MappedColumn)} + mapped_scalars = set(entry_type.__table__.columns.keys()) + mapped_scalars -= shared_columns + mapped_scalars -= set(entry_type.CHILD_HASH_COLUMNS.values()) + + assert mapped_scalars == set(identifier_type.promoted_scalar_field_names()) + + def test_atomic_attack_identifier_graph_persists_with_result_link() -> None: target = TargetIdentifier(class_name="Target", class_module="pyrit.prompt_target", model_name="model") scorer = ScorerIdentifier(class_name="Scorer", class_module="pyrit.score", scorer_type="true_false") From 907f49e3f561ee281e3f46a3062d2f51ef252c5c Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 13:45:49 -0700 Subject: [PATCH 17/24] add getters in memory interface --- pyrit/memory/memory_interface.py | 341 ++++++++++++++++++ .../test_interface_identifiers.py | 275 ++++++++++++++ 2 files changed, 616 insertions(+) create mode 100644 tests/unit/memory/memory_interface/test_interface_identifiers.py diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index acd7a604b8..6318279b3d 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -80,6 +80,7 @@ Model = TypeVar("Model") +IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) class MemoryInterface(abc.ABC): @@ -592,6 +593,346 @@ def _iter_identifier_dependencies(identifier: ComponentIdentifier) -> Iterator[C elif isinstance(child, list): yield from (item for item in child if isinstance(item, ComponentIdentifier)) + def _get_identifiers( + self, + *, + identifier_type: type[IdentifierModel], + entry_type: type[ComponentIdentifierEntry[Any]], + identifier_hashes: Sequence[str] | None, + filters: dict[str, Any], + ) -> Sequence[IdentifierModel]: + if identifier_hashes is not None and not identifier_hashes: + return [] + + conditions = [getattr(entry_type, name) == value for name, value in filters.items() if value is not None] + if identifier_hashes is not None: + entries = self._execute_batched_query( + entry_type, + batch_column=entry_type.hash, + batch_values=identifier_hashes, + other_conditions=conditions, + order_by=entry_type.hash, + ) + else: + entries = self._query_entries( + entry_type, + conditions=and_(*conditions) if conditions else None, + order_by=entry_type.hash, + ) + + identifiers: list[IdentifierModel] = [] + seen_hashes: set[str] = set() + for entry in sorted(entries, key=lambda item: item.hash): + if entry.hash in seen_hashes: + continue + if entry.identifier_json is None: + raise ValueError(f"Identifier row {entry.hash} in {entry_type.__tablename__} has no identifier JSON.") + identifier = identifier_type.model_validate(entry.identifier_json) + if identifier.hash != entry.hash: + raise ValueError( + f"Identifier row {entry.hash} in {entry_type.__tablename__} does not match its stored JSON hash." + ) + identifiers.append(identifier) + seen_hashes.add(entry.hash) + return identifiers + + def get_target_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + endpoint: str | None = None, + model_name: str | None = None, + underlying_model_name: str | None = None, + temperature: float | None = None, + top_p: float | None = None, + max_requests_per_minute: int | None = None, + supported_auth_modes: Sequence[str] | None = None, + ) -> Sequence[TargetIdentifier]: + """ + Retrieve target identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + endpoint (str | None): Target endpoint to match. + model_name (str | None): Target model name to match. + underlying_model_name (str | None): Underlying model name to match. + temperature (float | None): Temperature to match. + top_p (float | None): Top-p value to match. + max_requests_per_minute (int | None): Request limit to match. + supported_auth_modes (Sequence[str] | None): Ordered authentication modes to match. + + Returns: + Sequence[TargetIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=TargetIdentifier, + entry_type=TargetIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "endpoint": endpoint, + "model_name": model_name, + "underlying_model_name": underlying_model_name, + "temperature": temperature, + "top_p": top_p, + "max_requests_per_minute": max_requests_per_minute, + "supported_auth_modes": list(supported_auth_modes) if supported_auth_modes is not None else None, + }, + ) + + def get_converter_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + supported_input_types: Sequence[str] | None = None, + supported_output_types: Sequence[str] | None = None, + converter_target_hash: str | None = None, + sub_converter_hash: str | None = None, + ) -> Sequence[ConverterIdentifier]: + """ + Retrieve converter identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + supported_input_types (Sequence[str] | None): Ordered input types to match. + supported_output_types (Sequence[str] | None): Ordered output types to match. + converter_target_hash (str | None): Converter target hash to match. + sub_converter_hash (str | None): Nested converter hash to match. + + Returns: + Sequence[ConverterIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ConverterIdentifier, + entry_type=ConverterIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "supported_input_types": ( + list(supported_input_types) if supported_input_types is not None else None + ), + "supported_output_types": ( + list(supported_output_types) if supported_output_types is not None else None + ), + "converter_target_hash": converter_target_hash, + "sub_converter_hash": sub_converter_hash, + }, + ) + + def get_scorer_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + scorer_type: str | None = None, + score_aggregator: str | None = None, + prompt_target_hash: str | None = None, + ) -> Sequence[ScorerIdentifier]: + """ + Retrieve scorer identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + scorer_type (str | None): Scorer type to match. + score_aggregator (str | None): Score aggregator to match. + prompt_target_hash (str | None): Scorer target hash to match. + + Returns: + Sequence[ScorerIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ScorerIdentifier, + entry_type=ScorerIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "scorer_type": scorer_type, + "score_aggregator": score_aggregator, + "prompt_target_hash": prompt_target_hash, + }, + ) + + def get_scenario_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + version: int | None = None, + techniques: Sequence[str] | None = None, + datasets: Sequence[str] | None = None, + objective_target_hash: str | None = None, + objective_scorer_hash: str | None = None, + ) -> Sequence[ScenarioIdentifier]: + """ + Retrieve scenario identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + version (int | None): Scenario definition version to match. + techniques (Sequence[str] | None): Ordered technique names to match. + datasets (Sequence[str] | None): Ordered dataset names to match. + objective_target_hash (str | None): Objective target hash to match. + objective_scorer_hash (str | None): Objective scorer hash to match. + + Returns: + Sequence[ScenarioIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ScenarioIdentifier, + entry_type=ScenarioIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "version": version, + "techniques": list(techniques) if techniques is not None else None, + "datasets": list(datasets) if datasets is not None else None, + "objective_target_hash": objective_target_hash, + "objective_scorer_hash": objective_scorer_hash, + }, + ) + + def get_seed_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + value: str | None = None, + value_sha256: str | None = None, + data_type: str | None = None, + dataset_name: str | None = None, + is_general_technique: bool | None = None, + ) -> Sequence[SeedIdentifier]: + """ + Retrieve seed identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + value (str | None): Seed value to match. + value_sha256 (str | None): Seed value hash to match. + data_type (str | None): Seed data type to match. + dataset_name (str | None): Dataset name to match. + is_general_technique (bool | None): General-technique flag to match. + + Returns: + Sequence[SeedIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=SeedIdentifier, + entry_type=SeedIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "value": value, + "value_sha256": value_sha256, + "data_type": data_type, + "dataset_name": dataset_name, + "is_general_technique": is_general_technique, + }, + ) + + def get_attack_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + adversarial_system_prompt: str | None = None, + adversarial_seed_prompt: str | None = None, + objective_target_hash: str | None = None, + adversarial_chat_hash: str | None = None, + objective_scorer_hash: str | None = None, + ) -> Sequence[AttackIdentifier]: + """ + Retrieve attack identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + adversarial_system_prompt (str | None): Adversarial system prompt to match. + adversarial_seed_prompt (str | None): Adversarial seed prompt to match. + objective_target_hash (str | None): Objective target hash to match. + adversarial_chat_hash (str | None): Adversarial chat target hash to match. + objective_scorer_hash (str | None): Objective scorer hash to match. + + Returns: + Sequence[AttackIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AttackIdentifier, + entry_type=AttackIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "adversarial_system_prompt": adversarial_system_prompt, + "adversarial_seed_prompt": adversarial_seed_prompt, + "objective_target_hash": objective_target_hash, + "adversarial_chat_hash": adversarial_chat_hash, + "objective_scorer_hash": objective_scorer_hash, + }, + ) + + def get_attack_technique_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + attack_identifier_hash: str | None = None, + ) -> Sequence[AttackTechniqueIdentifier]: + """ + Retrieve attack technique identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + attack_identifier_hash (str | None): Attack identifier hash to match. + + Returns: + Sequence[AttackTechniqueIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AttackTechniqueIdentifier, + entry_type=AttackTechniqueIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "attack_identifier_hash": attack_identifier_hash, + }, + ) + + def get_atomic_attack_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + attack_technique_identifier_hash: str | None = None, + ) -> Sequence[AtomicAttackIdentifier]: + """ + Retrieve atomic attack identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + attack_technique_identifier_hash (str | None): Attack technique hash to match. + + Returns: + Sequence[AtomicAttackIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AtomicAttackIdentifier, + entry_type=AtomicAttackIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "attack_technique_identifier_hash": attack_technique_identifier_hash, + }, + ) + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ diff --git a/tests/unit/memory/memory_interface/test_interface_identifiers.py b/tests/unit/memory/memory_interface/test_interface_identifiers.py new file mode 100644 index 0000000000..345ad92995 --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_identifiers.py @@ -0,0 +1,275 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from contextlib import closing +from dataclasses import dataclass + +import pytest + +from pyrit.memory import MemoryInterface +from pyrit.memory.memory_models import TargetIdentifierEntry +from pyrit.models import ( + AtomicAttackIdentifier, + AttackIdentifier, + AttackTechniqueIdentifier, + Conversation, + ConverterIdentifier, + ScenarioIdentifier, + ScorerIdentifier, + SeedIdentifier, + TargetIdentifier, +) + + +@dataclass(frozen=True) +class IdentifierGraph: + objective_target: TargetIdentifier + adversarial_target: TargetIdentifier + nested_converter: ConverterIdentifier + converter: ConverterIdentifier + scorer: ScorerIdentifier + scenario: ScenarioIdentifier + technique_seed: SeedIdentifier + dataset_seed: SeedIdentifier + attack: AttackIdentifier + technique: AttackTechniqueIdentifier + atomic_attack: AtomicAttackIdentifier + + +@pytest.fixture +def identifier_graph(sqlite_instance: MemoryInterface) -> IdentifierGraph: + objective_target = TargetIdentifier( + class_name="ObjectiveTarget", + class_module="tests.unit.memory", + endpoint="https://objective.test", + model_name="objective-model", + temperature=0.5, + ) + adversarial_target = TargetIdentifier( + class_name="AdversarialTarget", + class_module="tests.unit.memory", + endpoint="https://adversarial.test", + model_name="adversarial-model", + ) + nested_converter = ConverterIdentifier( + class_name="NestedConverter", + class_module="tests.unit.memory", + supported_input_types=["text"], + supported_output_types=["text"], + converter_target=adversarial_target, + ) + converter = ConverterIdentifier( + class_name="CompositeConverter", + class_module="tests.unit.memory", + supported_input_types=["text"], + supported_output_types=["text"], + sub_converter=nested_converter, + ) + scorer = ScorerIdentifier( + class_name="TestScorer", + class_module="tests.unit.memory", + scorer_type="true_false", + score_aggregator="AND_", + prompt_target=adversarial_target, + ) + technique_seed = SeedIdentifier( + class_name="TechniqueSeed", + class_module="tests.unit.memory", + value="technique seed", + value_sha256="technique-sha", + data_type="text", + dataset_name="techniques", + is_general_technique=True, + ) + dataset_seed = SeedIdentifier( + class_name="DatasetSeed", + class_module="tests.unit.memory", + value="dataset seed", + value_sha256="dataset-sha", + data_type="text", + dataset_name="dataset-a", + is_general_technique=False, + ) + attack = AttackIdentifier( + class_name="TestAttack", + class_module="tests.unit.memory", + adversarial_system_prompt="system prompt", + adversarial_seed_prompt="seed prompt", + objective_target=objective_target, + adversarial_chat=adversarial_target, + objective_scorer=scorer, + request_converters=[converter], + ) + technique = AttackTechniqueIdentifier( + class_name="TestTechnique", + class_module="tests.unit.memory", + attack=attack, + technique_seeds=[technique_seed], + ) + atomic_attack = AtomicAttackIdentifier( + class_name="AtomicAttack", + class_module="tests.unit.memory", + attack_technique=technique, + seed_identifiers=[technique_seed, dataset_seed], + ) + scenario = ScenarioIdentifier( + class_name="TestScenario", + class_module="tests.unit.memory", + version=2, + techniques=["TestTechnique"], + datasets=["dataset-a"], + objective_target=objective_target, + objective_scorer=scorer, + ) + + with closing(sqlite_instance.get_session()) as session: + sqlite_instance._persist_identifier(session=session, identifier=atomic_attack) + sqlite_instance._persist_identifier(session=session, identifier=scenario) + session.commit() + + return IdentifierGraph( + objective_target=objective_target, + adversarial_target=adversarial_target, + nested_converter=nested_converter, + converter=converter, + scorer=scorer, + scenario=scenario, + technique_seed=technique_seed, + dataset_seed=dataset_seed, + attack=attack, + technique=technique, + atomic_attack=atomic_attack, + ) + + +def test_get_target_identifiers_by_hash_and_promoted_field(sqlite_instance: MemoryInterface) -> None: + target = TargetIdentifier( + class_name="TestTarget", + class_module="tests.unit.memory", + endpoint="https://example.test", + model_name="test-model", + supported_auth_modes=["api_key"], + ) + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id="identifier-query", target_identifier=target) + ) + + identifiers = sqlite_instance.get_target_identifiers( + identifier_hashes=[target.hash], + model_name="test-model", + supported_auth_modes=["api_key"], + ) + + assert identifiers == [target] + assert isinstance(identifiers[0], TargetIdentifier) + + +def test_get_identifiers_reconstructs_each_typed_graph( + sqlite_instance: MemoryInterface, identifier_graph: IdentifierGraph +) -> None: + queries = [ + ( + sqlite_instance.get_target_identifiers( + identifier_hashes=[identifier_graph.objective_target.hash], + endpoint="https://objective.test", + temperature=0.5, + ), + identifier_graph.objective_target, + ), + ( + sqlite_instance.get_converter_identifiers( + supported_input_types=["text"], + supported_output_types=["text"], + sub_converter_hash=identifier_graph.nested_converter.hash, + ), + identifier_graph.converter, + ), + ( + sqlite_instance.get_scorer_identifiers( + scorer_type="true_false", + score_aggregator="AND_", + prompt_target_hash=identifier_graph.adversarial_target.hash, + ), + identifier_graph.scorer, + ), + ( + sqlite_instance.get_scenario_identifiers( + version=2, + techniques=["TestTechnique"], + datasets=["dataset-a"], + objective_target_hash=identifier_graph.objective_target.hash, + objective_scorer_hash=identifier_graph.scorer.hash, + ), + identifier_graph.scenario, + ), + ( + sqlite_instance.get_seed_identifiers( + value_sha256="dataset-sha", + data_type="text", + dataset_name="dataset-a", + is_general_technique=False, + ), + identifier_graph.dataset_seed, + ), + ( + sqlite_instance.get_attack_identifiers( + adversarial_system_prompt="system prompt", + adversarial_seed_prompt="seed prompt", + objective_target_hash=identifier_graph.objective_target.hash, + adversarial_chat_hash=identifier_graph.adversarial_target.hash, + objective_scorer_hash=identifier_graph.scorer.hash, + ), + identifier_graph.attack, + ), + ( + sqlite_instance.get_attack_technique_identifiers( + attack_identifier_hash=identifier_graph.attack.hash, + ), + identifier_graph.technique, + ), + ( + sqlite_instance.get_atomic_attack_identifiers( + attack_technique_identifier_hash=identifier_graph.technique.hash, + ), + identifier_graph.atomic_attack, + ), + ] + + for identifiers, expected in queries: + assert identifiers == [expected] + assert type(identifiers[0]) is type(expected) + + +def test_get_identifiers_common_filters_and_result_semantics( + sqlite_instance: MemoryInterface, identifier_graph: IdentifierGraph +) -> None: + targets = sqlite_instance.get_target_identifiers() + assert [identifier.hash for identifier in targets] == sorted( + [identifier_graph.objective_target.hash, identifier_graph.adversarial_target.hash] + ) + + duplicate_hashes = [identifier_graph.objective_target.hash, identifier_graph.objective_target.hash] + assert sqlite_instance.get_target_identifiers(identifier_hashes=duplicate_hashes) == [ + identifier_graph.objective_target + ] + assert sqlite_instance.get_target_identifiers(identifier_hashes=[]) == [] + assert sqlite_instance.get_target_identifiers(class_name="missing") == [] + + +def test_get_identifiers_rejects_missing_identifier_json(sqlite_instance: MemoryInterface) -> None: + identifier_hash = "0" * 64 + sqlite_instance._insert_entry(TargetIdentifierEntry(hash=identifier_hash, identifier_json=None)) + + with pytest.raises(ValueError, match="has no identifier JSON"): + sqlite_instance.get_target_identifiers(identifier_hashes=[identifier_hash]) + + +def test_get_identifiers_rejects_hash_mismatch(sqlite_instance: MemoryInterface) -> None: + target = TargetIdentifier(class_name="Target", class_module="tests.unit.memory") + identifier_hash = "f" * 64 + sqlite_instance._insert_entry( + TargetIdentifierEntry(hash=identifier_hash, identifier_json=target.model_dump()) + ) + + with pytest.raises(ValueError, match="does not match its stored JSON hash"): + sqlite_instance.get_target_identifiers(identifier_hashes=[identifier_hash]) From cb579254051df231271c9c02355a8cdc75fd6845 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 14:22:48 -0700 Subject: [PATCH 18/24] fix cascade delete orphans --- pyrit/memory/memory_models.py | 2 ++ tests/unit/memory/test_memory_models.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 8aa30514ac..43536d17a8 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -591,6 +591,7 @@ class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): foreign_keys="TargetIdentifierChildEntry.parent_hash", order_by="TargetIdentifierChildEntry.position", back_populates="parent", + cascade="all, delete-orphan", ) @@ -725,6 +726,7 @@ class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): foreign_keys="ScorerIdentifierChildEntry.parent_hash", order_by="ScorerIdentifierChildEntry.position", back_populates="parent", + cascade="all, delete-orphan", ) diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 3c2c357346..16a8f898b0 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -245,6 +245,20 @@ def test_identifier_entry_maps_promoted_children_by_cardinality( assert set(entry_type.CHILD_HASH_COLUMNS) == singular_children +def test_identifier_child_relationships_delete_orphans() -> None: + for mapper in Base.registry.mappers: + entry_type = mapper.class_ + if not issubclass(entry_type, ComponentIdentifierEntry): + continue + + for field_name, spec in entry_type.CHILD_RELATIONSHIP_SPECS.items(): + child_relationship = mapper.relationships[spec.relationship_name] + assert child_relationship.uselist, f"{entry_type.__name__}.{field_name} must be a collection" + assert "delete-orphan" in child_relationship.cascade, ( + f"{entry_type.__name__}.{spec.relationship_name} must use cascade='all, delete-orphan'" + ) + + @pytest.mark.parametrize( ("identifier_type", "entry_type"), [ From e06e5ec03c4a10a8c5967c583145115e5d8023eb Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 14:51:05 -0700 Subject: [PATCH 19/24] ruff --- pyrit/memory/memory_interface.py | 4 +--- .../memory/memory_interface/test_interface_identifiers.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 6318279b3d..c1f021da06 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -712,9 +712,7 @@ def get_converter_identifiers( identifier_hashes=identifier_hashes, filters={ "class_name": class_name, - "supported_input_types": ( - list(supported_input_types) if supported_input_types is not None else None - ), + "supported_input_types": (list(supported_input_types) if supported_input_types is not None else None), "supported_output_types": ( list(supported_output_types) if supported_output_types is not None else None ), diff --git a/tests/unit/memory/memory_interface/test_interface_identifiers.py b/tests/unit/memory/memory_interface/test_interface_identifiers.py index 345ad92995..ff585fd982 100644 --- a/tests/unit/memory/memory_interface/test_interface_identifiers.py +++ b/tests/unit/memory/memory_interface/test_interface_identifiers.py @@ -267,9 +267,7 @@ def test_get_identifiers_rejects_missing_identifier_json(sqlite_instance: Memory def test_get_identifiers_rejects_hash_mismatch(sqlite_instance: MemoryInterface) -> None: target = TargetIdentifier(class_name="Target", class_module="tests.unit.memory") identifier_hash = "f" * 64 - sqlite_instance._insert_entry( - TargetIdentifierEntry(hash=identifier_hash, identifier_json=target.model_dump()) - ) + sqlite_instance._insert_entry(TargetIdentifierEntry(hash=identifier_hash, identifier_json=target.model_dump())) with pytest.raises(ValueError, match="does not match its stored JSON hash"): sqlite_instance.get_target_identifiers(identifier_hashes=[identifier_hash]) From 0d92c683bb66d0858ea882e85e960cc7a69fbd54 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 15:15:51 -0700 Subject: [PATCH 20/24] idempotent migration script --- .../e5f7a9c1b3d2_add_identifiers_tables.py | 226 +++++++++++++----- tests/unit/memory/test_migration.py | 84 +++++++ 2 files changed, 249 insertions(+), 61 deletions(-) diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py index 9d66a62a90..a3e611b1d3 100644 --- a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py @@ -19,6 +19,7 @@ import logging import uuid from collections.abc import Sequence # noqa: TC003 +from functools import partial from typing import TYPE_CHECKING, Any import sqlalchemy as sa @@ -72,6 +73,37 @@ def run_best_effort_backfill(*, bind: Any, name: str, backfill: Callable[[], Non logger.warning(f"{name} backfill failed; leaving new identifier links nullable", exc_info=True) +def _run_best_effort_row(*, bind: Any, description: str, operation: Callable[[], None]) -> bool: + """ + Run one backfill row in a savepoint and report whether it succeeded. + + Returns: + bool: Whether the row operation completed successfully. + """ + try: + with bind.begin_nested(): + operation() + except Exception: + logger.warning(description, exc_info=True) + return False + return True + + +def _insert_identifier_link( + *, + bind: Any, + insert: Callable[[dict[str, Any]], str | None], + identifier: dict[str, Any], + update_statement: Any, + update_values: dict[str, Any], +) -> None: + """Insert an identifier graph and update its domain-row link.""" + identifier_hash = insert(identifier) + if not identifier_hash: + return + bind.execute(update_statement, {**update_values, "hash": identifier_hash}) + + def load_identifier(raw_identifier: Any) -> dict[str, Any] | None: """ Load a retained identifier JSON value without importing domain models. @@ -361,14 +393,27 @@ def _insert_edges( child_column: str, child_hashes: Sequence[str], ) -> None: + select_statement = sa.text( + f'SELECT "{child_column}" FROM "{table}" ' + f'WHERE "{parent_column}" = :parent_hash AND position = :position' + ) statement = sa.text( - f'INSERT INTO "{table}" ({parent_column}, position, {child_column}) ' + f'INSERT INTO "{table}" ("{parent_column}", position, "{child_column}") ' f"VALUES (:parent_hash, :position, :child_hash)" ) for position, child_hash in enumerate(child_hashes): + parameters = {"parent_hash": parent_hash, "position": position} + existing_child_hash = self._bind.execute(select_statement, parameters).scalar_one_or_none() + if existing_child_hash == child_hash: + continue + if existing_child_hash is not None: + raise ValueError( + f"Conflicting {table} edge for parent {parent_hash!r} at position {position}: " + f"stored child {existing_child_hash!r}, retained child {child_hash!r}." + ) self._bind.execute( statement, - {"parent_hash": parent_hash, "position": position, "child_hash": child_hash}, + {**parameters, "child_hash": child_hash}, ) @staticmethod @@ -570,12 +615,16 @@ def upgrade() -> None: def downgrade() -> None: """Revert this schema upgrade.""" with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.drop_constraint("fk_attack_result_entries_atomic_attack_identifier_hash", type_="foreignkey") batch_op.drop_column("atomic_attack_identifier_hash") with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.drop_constraint("fk_scenario_result_entries_scenario_identifier_hash", type_="foreignkey") batch_op.drop_column("scenario_identifier_hash") with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.drop_constraint("fk_score_entries_scorer_identifier_hash", type_="foreignkey") batch_op.drop_column("scorer_identifier_hash") with op.batch_alter_table("Conversations") as batch_op: + batch_op.drop_constraint("fk_conversations_target_identifier_hash", type_="foreignkey") batch_op.drop_column("target_identifier_hash") op.drop_table("AtomicAttackSeedIdentifiers") @@ -688,13 +737,38 @@ def _create_attack_identifier_tables() -> None: ) +def _insert_converter_links( + *, + bind: Any, + inserter: IdentifierGraphInserter, + link_statement: Any, + prompt_id: Any, + stored_identifiers: Any, + pyrit_version: str | None, +) -> None: + """Insert converter graphs and their ordered links for one prompt row.""" + for position, identifier in enumerate(load_identifier_list(stored_identifiers)): + if identifier.get("pyrit_version") is None: + identifier = {**identifier, "pyrit_version": pyrit_version} + identifier_hash = inserter.insert_converter(identifier) + if identifier_hash: + bind.execute( + link_statement, + { + "prompt_memory_entry_id": prompt_id, + "position": position, + "converter_identifier_hash": identifier_hash, + }, + ) + + def _backfill_target_identifiers() -> None: """ Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set ``Conversations.target_identifier_hash``. For every ``Conversations`` row with a non-null ``target_identifier`` JSON, - reconstruct the ``TargetIdentifier`` (recomputing its content hash), insert the + load the retained ``TargetIdentifier`` shape and its stored hash, insert the deduped ``TargetIdentifiers`` row if absent -- recursing into any inner ``targets`` first so the child edge foreign keys resolve -- record the ``parent_hash -> child_hash`` edges, and point the conversation's @@ -704,7 +778,10 @@ def _backfill_target_identifiers() -> None: """ bind = op.get_bind() rows = bind.execute( - sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') + sa.text( + 'SELECT conversation_id, target_identifier FROM "Conversations" ' + "WHERE target_identifier IS NOT NULL ORDER BY conversation_id" + ) ).fetchall() update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') @@ -716,14 +793,23 @@ def _backfill_target_identifiers() -> None: if identifier is None: skipped += 1 continue - try: - identifier_hash = inserter.insert_target(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "cid": conversation_id}) - linked += 1 - except Exception: + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_target, + identifier=identifier, + update_statement=update_stmt, + update_values={"cid": conversation_id}, + ) + if _run_best_effort_row( + bind=bind, + description=f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", + operation=operation, + ): + linked += 1 + else: skipped += 1 - logger.warning(f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", exc_info=True) + inserter = IdentifierGraphInserter(bind=bind) if linked or skipped: logger.info(f"TargetIdentifiers backfill linked {linked} conversation(s); skipped {skipped}.") @@ -733,7 +819,10 @@ def _backfill_scorer_identifiers() -> None: """Backfill scorer rows and score foreign keys from retained JSON.""" bind = op.get_bind() score_rows = bind.execute( - sa.text('SELECT id, scorer_class_identifier FROM "ScoreEntries" WHERE scorer_class_identifier IS NOT NULL') + sa.text( + 'SELECT id, scorer_class_identifier FROM "ScoreEntries" ' + "WHERE scorer_class_identifier IS NOT NULL ORDER BY id" + ) ).fetchall() score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') inserter = IdentifierGraphInserter(bind=bind) @@ -743,16 +832,21 @@ def _backfill_scorer_identifiers() -> None: if identifier is None: skipped += 1 continue - try: - identifier_hash = inserter.insert_scorer(identifier) - if identifier_hash: - bind.execute(score_update, {"hash": identifier_hash, "id": score_id}) - except Exception: + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_scorer, + identifier=identifier, + update_statement=score_update, + update_values={"id": score_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", + operation=operation, + ): skipped += 1 - logger.warning( - f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", - exc_info=True, - ) + inserter = IdentifierGraphInserter(bind=bind) if skipped: logger.warning(f"ScorerIdentifiers backfill skipped {skipped} score row(s)") @@ -761,7 +855,10 @@ def _backfill_scenario_identifiers() -> None: """Backfill scenario rows and result foreign keys from retained JSON.""" bind = op.get_bind() result_rows = bind.execute( - sa.text('SELECT id, scenario_identifier FROM "ScenarioResultEntries" WHERE scenario_identifier IS NOT NULL') + sa.text( + 'SELECT id, scenario_identifier FROM "ScenarioResultEntries" ' + "WHERE scenario_identifier IS NOT NULL ORDER BY id" + ) ).fetchall() update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') inserter = IdentifierGraphInserter(bind=bind) @@ -771,16 +868,21 @@ def _backfill_scenario_identifiers() -> None: if identifier is None: skipped += 1 continue - try: - identifier_hash = inserter.insert_scenario(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) - except Exception: + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_scenario, + identifier=identifier, + update_statement=update_stmt, + update_values={"id": result_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", + operation=operation, + ): skipped += 1 - logger.warning( - f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", - exc_info=True, - ) + inserter = IdentifierGraphInserter(bind=bind) if skipped: logger.warning(f"ScenarioIdentifiers backfill skipped {skipped} scenario result row(s)") @@ -791,7 +893,7 @@ def _backfill_converter_identifiers() -> None: prompt_rows = bind.execute( sa.text( 'SELECT id, converter_identifiers, pyrit_version FROM "PromptMemoryEntries" ' - "WHERE converter_identifiers IS NOT NULL" + "WHERE converter_identifiers IS NOT NULL ORDER BY id" ) ).fetchall() link_insert = sa.text( @@ -802,26 +904,22 @@ def _backfill_converter_identifiers() -> None: inserter = IdentifierGraphInserter(bind=bind) skipped = 0 for prompt_id, stored_identifiers, pyrit_version in prompt_rows: - try: - for position, identifier in enumerate(load_identifier_list(stored_identifiers)): - if identifier.get("pyrit_version") is None: - identifier = {**identifier, "pyrit_version": pyrit_version} - identifier_hash = inserter.insert_converter(identifier) - if identifier_hash: - bind.execute( - link_insert, - { - "prompt_memory_entry_id": prompt_id, - "position": position, - "converter_identifier_hash": identifier_hash, - }, - ) - except Exception: + operation = partial( + _insert_converter_links, + bind=bind, + inserter=inserter, + link_statement=link_insert, + prompt_id=prompt_id, + stored_identifiers=stored_identifiers, + pyrit_version=pyrit_version, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", + operation=operation, + ): skipped += 1 - logger.warning( - f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", - exc_info=True, - ) + inserter = IdentifierGraphInserter(bind=bind) if skipped: logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") @@ -831,7 +929,8 @@ def _backfill_attack_identifiers() -> None: bind = op.get_bind() result_rows = bind.execute( sa.text( - 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" WHERE atomic_attack_identifier IS NOT NULL' + 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" ' + "WHERE atomic_attack_identifier IS NOT NULL ORDER BY id" ) ).fetchall() inserter = IdentifierGraphInserter(bind=bind) @@ -842,15 +941,20 @@ def _backfill_attack_identifiers() -> None: if identifier is None: skipped += 1 continue - try: - identifier_hash = inserter.insert_atomic_attack(identifier) - if identifier_hash: - bind.execute(update_stmt, {"hash": identifier_hash, "id": result_id}) - except Exception: + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_atomic_attack, + identifier=identifier, + update_statement=update_stmt, + update_values={"id": result_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"Attack identifier backfill could not reconstruct result {result_id}", + operation=operation, + ): skipped += 1 - logger.warning( - f"Attack identifier backfill could not reconstruct result {result_id}", - exc_info=True, - ) + inserter = IdentifierGraphInserter(bind=bind) if skipped: logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 3b89efd51c..116ec7f7ca 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -921,6 +921,90 @@ def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json( engine.dispose() +def test_identifier_migration_reuses_edges_and_rolls_back_conflicting_row(): + """Repeated graphs reuse edges while a conflicting row is isolated in its savepoint.""" + child_hash = "a" * 64 + conflicting_child_hash = "b" * 64 + parent_hash = "c" * 64 + healthy_hash = "d" * 64 + child = {"hash": child_hash, "class_name": "Child", "class_module": "test"} + parent = { + "hash": parent_hash, + "class_name": "Parent", + "class_module": "test", + "children": {"targets": [child]}, + } + conflicting_parent = { + **parent, + "children": { + "targets": [{"hash": conflicting_child_hash, "class_name": "OtherChild", "class_module": "test"}] + }, + } + healthy = {"hash": healthy_hash, "class_name": "Healthy", "class_module": "test"} + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-row-savepoints.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "d4e6f8a0b2c4") + insert_statement = text( + 'INSERT INTO "Conversations" (conversation_id, target_identifier) VALUES (:id, :identifier)' + ) + for conversation_id, identifier in ( + ("01-original", parent), + ("02-duplicate", parent), + ("03-conflict", conflicting_parent), + ("04-healthy", healthy), + ): + connection.execute( + insert_statement, + {"id": conversation_id, "identifier": json.dumps(identifier)}, + ) + + command.upgrade(config, "e5f7a9c1b3d2") + + links = connection.execute( + text( + 'SELECT conversation_id, target_identifier_hash FROM "Conversations" ' + "ORDER BY conversation_id" + ) + ).fetchall() + target_hashes = set(connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars()) + edges = connection.execute( + text('SELECT parent_hash, position, child_hash FROM "TargetIdentifierChildren"') + ).fetchall() + + assert links == [ + ("01-original", parent_hash), + ("02-duplicate", parent_hash), + ("03-conflict", None), + ("04-healthy", healthy_hash), + ] + assert target_hashes == {child_hash, parent_hash, healthy_hash} + assert edges == [(parent_hash, 0, child_hash)] + finally: + engine.dispose() + + +def test_identifier_migration_downgrade_drops_link_constraints_and_columns(): + """Downgrade removes identifier foreign keys before removing their columns.""" + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-downgrade.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "e5f7a9c1b3d2") + command.downgrade(config, "d4e6f8a0b2c4") + + assert "target_identifier_hash" not in { + column["name"] for column in inspect(connection).get_columns("Conversations") + } + assert "TargetIdentifiers" not in set(inspect(connection).get_table_names()) + finally: + engine.dispose() + + def test_identifier_migrations_do_not_import_domain_models(): """Frozen identifier migrations operate on retained JSON rather than current domain models.""" versions_dir = Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic" / "versions" From 16be52a0a857d8cf09e18e59f3bb4d9839e02881 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 15:40:08 -0700 Subject: [PATCH 21/24] identifier list filters --- pyrit/memory/memory_interface.py | 25 +++++++++++++----- .../test_interface_identifiers.py | 26 ++++++++++--------- 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index c1f021da06..1d9db58b4a 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -604,7 +604,12 @@ def _get_identifiers( if identifier_hashes is not None and not identifier_hashes: return [] - conditions = [getattr(entry_type, name) == value for name, value in filters.items() if value is not None] + list_filters = {name: value for name, value in filters.items() if isinstance(value, list)} + conditions = [ + getattr(entry_type, name) == value + for name, value in filters.items() + if value is not None and name not in list_filters + ] if identifier_hashes is not None: entries = self._execute_batched_query( entry_type, @@ -620,6 +625,14 @@ def _get_identifiers( order_by=entry_type.hash, ) + entries = [ + entry + for entry in entries + if all( + getattr(entry, name) is not None and sorted(getattr(entry, name)) == sorted(value) + for name, value in list_filters.items() + ) + ] identifiers: list[IdentifierModel] = [] seen_hashes: set[str] = set() for entry in sorted(entries, key=lambda item: item.hash): @@ -661,7 +674,7 @@ def get_target_identifiers( temperature (float | None): Temperature to match. top_p (float | None): Top-p value to match. max_requests_per_minute (int | None): Request limit to match. - supported_auth_modes (Sequence[str] | None): Ordered authentication modes to match. + supported_auth_modes (Sequence[str] | None): Authentication modes to match exactly, in any order. Returns: Sequence[TargetIdentifier]: Matching identifiers ordered by content hash. @@ -698,8 +711,8 @@ def get_converter_identifiers( Args: identifier_hashes (Sequence[str] | None): Content hashes to include. class_name (str | None): Component class name to match. - supported_input_types (Sequence[str] | None): Ordered input types to match. - supported_output_types (Sequence[str] | None): Ordered output types to match. + supported_input_types (Sequence[str] | None): Input types to match exactly, in any order. + supported_output_types (Sequence[str] | None): Output types to match exactly, in any order. converter_target_hash (str | None): Converter target hash to match. sub_converter_hash (str | None): Nested converter hash to match. @@ -773,8 +786,8 @@ def get_scenario_identifiers( identifier_hashes (Sequence[str] | None): Content hashes to include. class_name (str | None): Component class name to match. version (int | None): Scenario definition version to match. - techniques (Sequence[str] | None): Ordered technique names to match. - datasets (Sequence[str] | None): Ordered dataset names to match. + techniques (Sequence[str] | None): Technique names to match exactly, in any order. + datasets (Sequence[str] | None): Dataset names to match exactly, in any order. objective_target_hash (str | None): Objective target hash to match. objective_scorer_hash (str | None): Objective scorer hash to match. diff --git a/tests/unit/memory/memory_interface/test_interface_identifiers.py b/tests/unit/memory/memory_interface/test_interface_identifiers.py index ff585fd982..8f638655e7 100644 --- a/tests/unit/memory/memory_interface/test_interface_identifiers.py +++ b/tests/unit/memory/memory_interface/test_interface_identifiers.py @@ -54,15 +54,15 @@ def identifier_graph(sqlite_instance: MemoryInterface) -> IdentifierGraph: nested_converter = ConverterIdentifier( class_name="NestedConverter", class_module="tests.unit.memory", - supported_input_types=["text"], - supported_output_types=["text"], + supported_input_types=["text", "image_path"], + supported_output_types=["text", "audio_path"], converter_target=adversarial_target, ) converter = ConverterIdentifier( class_name="CompositeConverter", class_module="tests.unit.memory", - supported_input_types=["text"], - supported_output_types=["text"], + supported_input_types=["text", "image_path"], + supported_output_types=["text", "audio_path"], sub_converter=nested_converter, ) scorer = ScorerIdentifier( @@ -116,8 +116,8 @@ def identifier_graph(sqlite_instance: MemoryInterface) -> IdentifierGraph: class_name="TestScenario", class_module="tests.unit.memory", version=2, - techniques=["TestTechnique"], - datasets=["dataset-a"], + techniques=["TestTechnique", "OtherTechnique"], + datasets=["dataset-a", "dataset-b"], objective_target=objective_target, objective_scorer=scorer, ) @@ -148,7 +148,7 @@ def test_get_target_identifiers_by_hash_and_promoted_field(sqlite_instance: Memo class_module="tests.unit.memory", endpoint="https://example.test", model_name="test-model", - supported_auth_modes=["api_key"], + supported_auth_modes=["api_key", "identity"], ) sqlite_instance.add_conversation_to_memory( conversation=Conversation(conversation_id="identifier-query", target_identifier=target) @@ -157,11 +157,13 @@ def test_get_target_identifiers_by_hash_and_promoted_field(sqlite_instance: Memo identifiers = sqlite_instance.get_target_identifiers( identifier_hashes=[target.hash], model_name="test-model", - supported_auth_modes=["api_key"], + supported_auth_modes=["identity", "api_key"], ) assert identifiers == [target] assert isinstance(identifiers[0], TargetIdentifier) + assert sqlite_instance.get_target_identifiers(supported_auth_modes=["api_key"]) == [] + assert sqlite_instance.get_target_identifiers(supported_auth_modes=["API_KEY", "identity"]) == [] def test_get_identifiers_reconstructs_each_typed_graph( @@ -178,8 +180,8 @@ def test_get_identifiers_reconstructs_each_typed_graph( ), ( sqlite_instance.get_converter_identifiers( - supported_input_types=["text"], - supported_output_types=["text"], + supported_input_types=["image_path", "text"], + supported_output_types=["audio_path", "text"], sub_converter_hash=identifier_graph.nested_converter.hash, ), identifier_graph.converter, @@ -195,8 +197,8 @@ def test_get_identifiers_reconstructs_each_typed_graph( ( sqlite_instance.get_scenario_identifiers( version=2, - techniques=["TestTechnique"], - datasets=["dataset-a"], + techniques=["OtherTechnique", "TestTechnique"], + datasets=["dataset-b", "dataset-a"], objective_target_hash=identifier_graph.objective_target.hash, objective_scorer_hash=identifier_graph.scorer.hash, ), From 58a7847876a88ce25e3ec09930028c7d8da89165 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 15:40:23 -0700 Subject: [PATCH 22/24] identifer entry init_subclass checks --- pyrit/memory/memory_models.py | 29 ++++++++++++++++++++++++- tests/unit/memory/test_memory_models.py | 17 +++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 43536d17a8..2d0054038c 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -8,7 +8,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, ClassVar, Generic, Literal, TypeVar, cast +from typing import Any, ClassVar, Generic, Literal, TypeVar, cast, get_args, get_origin from pydantic import BaseModel, ConfigDict from sqlalchemy import ( @@ -477,6 +477,33 @@ class ComponentIdentifierEntry(DomainBackedEntry[T]): #: compatibility with existing databases. pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Validate that concrete entries persist every promoted scalar field. + + Raises: + TypeError: If the entry omits its identifier type or a mapped scalar column. + """ + super().__init_subclass__(**kwargs) + if cls.__dict__.get("__abstract__", False): + return + + identifier_type = next( + ( + get_args(base)[0] + for base in getattr(cls, "__orig_bases__", ()) + if get_origin(base) is ComponentIdentifierEntry + ), + None, + ) + if not isinstance(identifier_type, type) or not issubclass(identifier_type, ComponentIdentifier): + raise TypeError(f"{cls.__name__} must declare its ComponentIdentifier domain model type.") + + missing_columns = set(identifier_type.promoted_scalar_field_names()) - set(cls.__table__.columns.keys()) + if missing_columns: + names = ", ".join(sorted(missing_columns)) + raise TypeError(f"{cls.__name__} has no mapped column for promoted scalar field(s): {names}.") + @classmethod def from_domain_model(cls, domain_model: T) -> Self: """ diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 16a8f898b0..38710499f2 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -284,6 +284,23 @@ def test_identifier_entry_maps_promoted_scalars_to_columns( assert mapped_scalars == set(identifier_type.promoted_scalar_field_names()) +def test_identifier_entry_rejects_missing_promoted_scalar_column() -> None: + class IdentifierWithPromotedScalar(ComponentIdentifier): + promoted_value: str | None = None + + table_name = "IncompleteIdentifierEntries" + try: + with pytest.raises(TypeError, match="has no mapped column for promoted scalar field.*promoted_value"): + + class IncompleteIdentifierEntry(ComponentIdentifierEntry[IdentifierWithPromotedScalar]): + __tablename__ = table_name + __table_args__ = {"extend_existing": True} + finally: + table = Base.metadata.tables.get(table_name) + if table is not None: + Base.metadata.remove(table) + + def test_atomic_attack_identifier_graph_persists_with_result_link() -> None: target = TargetIdentifier(class_name="Target", class_module="pyrit.prompt_target", model_name="model") scorer = ScorerIdentifier(class_name="Scorer", class_module="pyrit.score", scorer_type="true_false") From 77b0307d77f2c05ef7775b7280f6bf336bbd05e3 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 15:41:25 -0700 Subject: [PATCH 23/24] ruff --- .../versions/e5f7a9c1b3d2_add_identifiers_tables.py | 3 +-- tests/unit/memory/test_migration.py | 9 ++------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py index a3e611b1d3..1cd6ab1e2c 100644 --- a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py @@ -394,8 +394,7 @@ def _insert_edges( child_hashes: Sequence[str], ) -> None: select_statement = sa.text( - f'SELECT "{child_column}" FROM "{table}" ' - f'WHERE "{parent_column}" = :parent_hash AND position = :position' + f'SELECT "{child_column}" FROM "{table}" WHERE "{parent_column}" = :parent_hash AND position = :position' ) statement = sa.text( f'INSERT INTO "{table}" ("{parent_column}", position, "{child_column}") ' diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 116ec7f7ca..a0cc626a54 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -936,9 +936,7 @@ def test_identifier_migration_reuses_edges_and_rolls_back_conflicting_row(): } conflicting_parent = { **parent, - "children": { - "targets": [{"hash": conflicting_child_hash, "class_name": "OtherChild", "class_module": "test"}] - }, + "children": {"targets": [{"hash": conflicting_child_hash, "class_name": "OtherChild", "class_module": "test"}]}, } healthy = {"hash": healthy_hash, "class_name": "Healthy", "class_module": "test"} @@ -965,10 +963,7 @@ def test_identifier_migration_reuses_edges_and_rolls_back_conflicting_row(): command.upgrade(config, "e5f7a9c1b3d2") links = connection.execute( - text( - 'SELECT conversation_id, target_identifier_hash FROM "Conversations" ' - "ORDER BY conversation_id" - ) + text('SELECT conversation_id, target_identifier_hash FROM "Conversations" ORDER BY conversation_id') ).fetchall() target_hashes = set(connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars()) edges = connection.execute( From 246ea9d1b81342841dd3b47634b362fc5f8c1f29 Mon Sep 17 00:00:00 2001 From: behnam Date: Tue, 14 Jul 2026 15:58:15 -0700 Subject: [PATCH 24/24] no recursive ORM entry build --- pyrit/memory/memory_interface.py | 11 +++---- pyrit/memory/memory_models.py | 42 +++++++------------------ tests/unit/memory/test_memory_models.py | 10 +++--- 3 files changed, 20 insertions(+), 43 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 1d9db58b4a..98f96a75f1 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -527,11 +527,10 @@ def _persist_target_identifier(cls, *, session: Any, target_identifier: TargetId """ Persist ``target_identifier`` and its inner targets as content-addressed rows. - The complete ORM graph is constructed from the domain model and merged into the - caller's session. SQLAlchemy reconciles shared child targets by primary key and - cascades the merge through ordered child edges. Identifier rows are immutable - and keyed by their content hash, so an identical target reused across many - conversations maps to a single row. + Dependencies are persisted before the target row, whose ordered child edges + reference those rows by content hash. Identifier rows are immutable and keyed + by their content hash, so an identical target reused across many conversations + maps to a single row. If the row already exists it was fully persisted before (children and edges included, since rows are immutable), so this returns early. Otherwise the row and @@ -556,7 +555,7 @@ def _persist_identifier(cls, *, session: Any, identifier: ComponentIdentifier) - try: with session.begin_nested(): - session.merge(entry_type.from_domain_model(identifier)) + session.add(entry_type.from_domain_model(identifier)) session.flush() except IntegrityError: with session.no_autoflush: diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 2d0054038c..59a93d5a4f 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -8,7 +8,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, ClassVar, Generic, Literal, TypeVar, cast, get_args, get_origin +from typing import Any, ClassVar, Generic, Literal, TypeVar, get_args, get_origin from pydantic import BaseModel, ConfigDict from sqlalchemy import ( @@ -435,8 +435,7 @@ class _ChildRelationshipSpec: relationship_name: str edge_factory: Callable[[], Any] - edge_child_attr: str | None = None - edge_child_hash_attr: str | None = None + edge_child_hash_attr: str edge_position_attr: str | None = "position" @@ -515,21 +514,6 @@ def from_domain_model(cls, domain_model: T) -> Self: Returns: Self: A new, unsaved row. """ - return cls._from_domain_model_recursive(domain_model=domain_model, seen={}) - - @classmethod - def _from_domain_model_recursive(cls, *, domain_model: T, seen: dict[str, Self]) -> Self: - existing = seen.get(domain_model.hash) - if existing is not None: - return existing - - entry = cls._from_domain_model_shallow(domain_model=domain_model) - seen[domain_model.hash] = entry - cls._attach_child_relationship_rows(entry=entry, domain_model=domain_model, seen=seen) - return entry - - @classmethod - def _from_domain_model_shallow(cls, *, domain_model: T) -> Self: entry = cls( hash=domain_model.hash, class_name=domain_model.class_name, @@ -540,6 +524,7 @@ def _from_domain_model_shallow(cls, *, domain_model: T) -> Self: for name, value in domain_model.promoted_scalar_values().items(): setattr(entry, name, value) # each promoted scalar → its mapped column cls._populate_child_hashes(entry=entry, domain_model=domain_model) + cls._attach_child_relationship_rows(entry=entry, domain_model=domain_model) return entry @classmethod @@ -549,7 +534,7 @@ def _populate_child_hashes(cls, *, entry: Self, domain_model: T) -> None: setattr(entry, hash_column, child.hash if child is not None else None) @classmethod - def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: dict[str, Self]) -> None: + def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T) -> None: for field_name in domain_model.promoted_child_field_names(): spec = cls.CHILD_RELATIONSHIP_SPECS.get(field_name) if spec is None: @@ -564,12 +549,7 @@ def _attach_child_relationship_rows(cls, *, entry: Self, domain_model: T, seen: if not isinstance(child_identifier, ComponentIdentifier): continue edge_row = spec.edge_factory() - if spec.edge_child_hash_attr is not None: - setattr(edge_row, spec.edge_child_hash_attr, child_identifier.hash) - elif spec.edge_child_attr is not None: - typed_child_identifier = cast("T", child_identifier) - child_entry = cls._from_domain_model_recursive(domain_model=typed_child_identifier, seen=seen) - setattr(edge_row, spec.edge_child_attr, child_entry) + setattr(edge_row, spec.edge_child_hash_attr, child_identifier.hash) if spec.edge_position_attr: setattr(edge_row, spec.edge_position_attr, position) edge_rows.append(edge_row) @@ -596,7 +576,7 @@ class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): "targets": _ChildRelationshipSpec( relationship_name="targets", edge_factory=lambda: TargetIdentifierChildEntry(), - edge_child_attr="child", + edge_child_hash_attr="child_hash", edge_position_attr="position", ) } @@ -636,10 +616,10 @@ class TargetIdentifierChildEntry(Base): ``TargetIdentifierEntry.identifier_json`` remains the source of truth for reconstruction (inner targets are stored inline there too). - Constructed as part of the complete graph returned by - ``TargetIdentifierEntry.from_domain_model`` and persisted by - ``MemoryInterface._persist_target_identifier``. It is a plain ``Base`` row rather - than a ``DomainBackedEntry`` because an edge has no standalone domain model. + Constructed with its child hash by ``TargetIdentifierEntry.from_domain_model`` + after ``MemoryInterface._persist_target_identifier`` has persisted the child row. + It is a plain ``Base`` row rather than a ``DomainBackedEntry`` because an edge has + no standalone domain model. """ __tablename__ = "TargetIdentifierChildren" @@ -731,7 +711,7 @@ class ScorerIdentifierEntry(ComponentIdentifierEntry[ScorerIdentifier]): "sub_scorers": _ChildRelationshipSpec( relationship_name="sub_scorers", edge_factory=lambda: ScorerIdentifierChildEntry(), - edge_child_attr="child", + edge_child_hash_attr="child_hash", edge_position_attr="position", ) } diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index 38710499f2..60ad1f0eb4 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -187,7 +187,7 @@ def test_load_identifier_injects_pyrit_version(): assert loaded.pyrit_version == "9.9.9" -def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): +def test_scorer_identifier_entry_constructs_hash_only_sub_scorer_edges(): leaf = ScorerIdentifier( class_name="LeafScorer", class_module="pyrit.score", @@ -209,11 +209,8 @@ def test_scorer_identifier_entry_constructs_full_sub_scorer_graph(): entry = ScorerIdentifierEntry.from_domain_model(domain_model=root) assert [edge.position for edge in entry.sub_scorers] == [0, 1] - assert entry.sub_scorers[0].child.hash == nested.hash - assert entry.sub_scorers[1].child.hash == leaf.hash - nested_entry = entry.sub_scorers[0].child - assert len(nested_entry.sub_scorers) == 1 - assert nested_entry.sub_scorers[0].child is entry.sub_scorers[1].child + assert [edge.child_hash for edge in entry.sub_scorers] == [nested.hash, leaf.hash] + assert all(edge.child is None for edge in entry.sub_scorers) @pytest.mark.parametrize( @@ -243,6 +240,7 @@ def test_identifier_entry_maps_promoted_children_by_cardinality( assert set(entry_type.CHILD_RELATIONSHIP_SPECS) == collection_children assert set(entry_type.CHILD_HASH_COLUMNS) == singular_children + assert all(spec.edge_child_hash_attr for spec in entry_type.CHILD_RELATIONSHIP_SPECS.values()) def test_identifier_child_relationships_delete_orphans() -> None: