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..669a64716c --- /dev/null +++ b/pyrit/memory/alembic/versions/a6c8e0f2b4d6_add_scorer_identifiers_table.py @@ -0,0 +1,189 @@ +# 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/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..859ef211c0 --- /dev/null +++ b/pyrit/memory/alembic/versions/b7d9f1a3c5e7_add_scenario_identifiers_table.py @@ -0,0 +1,215 @@ +# 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/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/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/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/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 9e2bf380af..8a70e65d98 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -7,28 +7,38 @@ 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 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: from pyrit.memory.memory_embedding import MemoryEmbedding from pyrit.memory.memory_models import ( + AtomicAttackIdentifierEntry, + AttackIdentifierEntry, AttackResultEntry, + AttackTechniqueIdentifierEntry, Base, + ComponentIdentifierEntry, ConversationEntry, + ConverterIdentifierEntry, EmbeddingDataEntry, + PromptConverterIdentifierEntry, PromptMemoryEntry, + ScenarioIdentifierEntry, ScenarioResultEntry, ScoreEntry, + ScorerIdentifierEntry, SeedEntry, + SeedIdentifierEntry, + TargetIdentifierEntry, ) from pyrit.memory.storage import ( DataTypeSerializer, @@ -37,19 +47,28 @@ set_seed_sha256_async, ) from pyrit.models import ( + AtomicAttackIdentifier, + AttackIdentifier, AttackResult, + AttackTechniqueIdentifier, + ComponentIdentifier, Conversation, ConversationStats, + ConverterIdentifier, IdentifierFilter, IdentifierType, Message, MessagePiece, + ScenarioIdentifier, ScenarioResult, Score, + ScorerIdentifier, Seed, SeedDataset, SeedGroup, + SeedIdentifier, SeedType, + TargetIdentifier, group_conversation_message_pieces_by_sequence, sort_message_pieces, ) @@ -410,6 +429,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: """ @@ -459,6 +499,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 +523,75 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}") raise + @classmethod + def _persist_target_identifier(cls, *, session: Any, target_identifier: TargetIdentifier) -> None: + """ + 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. + + 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. + """ + 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(): + session.merge(entry_type.from_domain_model(identifier)) + session.flush() + except IntegrityError: + pass + + @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): + return ConverterIdentifierEntry + 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: """ @@ -804,6 +920,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: @@ -815,7 +934,26 @@ 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 + + @classmethod + def _persist_scorer_identifier(cls, *, session: Any, scorer_identifier: ScorerIdentifier) -> None: + """Persist a complete scorer graph and its target dependencies.""" + cls._persist_identifier(session=session, identifier=scorer_identifier) def get_scores( self, @@ -1771,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: @@ -2127,10 +2273,28 @@ 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 + + @classmethod + def _persist_scenario_identifier(cls, *, session: Any, scenario_identifier: ScenarioIdentifier) -> None: + """Persist a scenario identifier and its target and scorer dependencies.""" + 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 a7cfdbf552..cdd7948985 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -4,15 +4,18 @@ import json import logging import uuid -from collections.abc import Sequence +from abc import abstractmethod +from collections.abc import Callable, Sequence +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Literal +from typing import Any, ClassVar, Generic, Literal, TypeVar, cast from pydantic import BaseModel, ConfigDict from sqlalchemy import ( ARRAY, INTEGER, JSON, + Boolean, DateTime, Float, ForeignKey, @@ -29,19 +32,24 @@ relationship, ) from sqlalchemy.types import Uuid +from typing_extensions import Self import pyrit from pyrit.common.utils import to_sha256 from pyrit.models import ( SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY, AtomicAttackEvaluationIdentifier, + AtomicAttackIdentifier, + AttackIdentifier, AttackOutcome, AttackResult, + AttackTechniqueIdentifier, ChatMessageRole, ComponentIdentifier, Conversation, ConversationReference, ConversationType, + ConverterIdentifier, EvaluationIdentifier, MessagePiece, PromptDataType, @@ -51,11 +59,14 @@ ScenarioRunState, Score, ScorerEvaluationIdentifier, + ScorerIdentifier, Seed, + SeedIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, SeedType, + TargetIdentifier, ) logger = logging.getLogger(__name__) @@ -278,6 +289,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: """ @@ -353,6 +371,617 @@ 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) + + +@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 | None = None + edge_child_hash_attr: str | None = None + edge_position_attr: str | None = "position" + + +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 + + #: 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. + 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) + + @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 + 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(): + 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 + 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) + if spec.edge_position_attr: + setattr(edge_row, spec.edge_position_attr, position) + edge_rows.append(edge_row) + + +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} + + 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) + 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) + + #: 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): + """ + 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). + + 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" + __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 + ) + + #: 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 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.""" + + __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", + ) + } + 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) + 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", + ) + + +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 ScenarioIdentifierEntry(ComponentIdentifierEntry[ScenarioIdentifier]): + """Content-addressed store of ``ScenarioIdentifier`` projections.""" + + __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) + 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], + ) + + +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``. @@ -362,6 +991,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 +1002,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 +1021,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: @@ -443,6 +1082,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 @@ -474,6 +1116,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 @@ -858,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 @@ -911,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: """ @@ -924,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 @@ -1135,6 +1789,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") @@ -1177,6 +1838,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/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/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]: """ diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index b3860519e3..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() @@ -679,6 +731,105 @@ 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-inner-a", target_identifier=inner_a) + ) + 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. 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/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..759d5d9424 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -2,37 +2,62 @@ # 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 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, EmbeddingDataEntry, EmbeddingMessageWithSimilarity, PromptMemoryEntry, + ScenarioIdentifierEntry, ScenarioResultEntry, ScoreEntry, + ScorerIdentifierEntry, SeedEntry, + SeedIdentifierEntry, + TargetIdentifierEntry, UTCDateTime, _load_identifier, ) from pyrit.models import ( AtomicAttackIdentifier, + AttackIdentifier, AttackOutcome, AttackResult, + AttackTechniqueIdentifier, ComponentIdentifier, ConversationReference, ConversationType, + ConverterIdentifier, MessagePiece, + ScenarioIdentifier, ScenarioResult, Score, + ScorerIdentifier, + SeedIdentifier, SeedObjective, SeedPrompt, SeedSimulatedConversation, + TargetIdentifier, ) from unit.mocks import make_scenario_result @@ -162,6 +187,140 @@ 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 + + +@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_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 + + +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 015642ca89..17988a8f4f 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,352 @@ 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() + + +# ============================================================================= +# 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) +# ============================================================================= + + +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() + + +# ============================================================================= +# 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"}