diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py new file mode 100644 index 0000000000..1cd6ab1e2c --- /dev/null +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_identifiers_tables.py @@ -0,0 +1,959 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Persist component identifiers as content-addressed rows. + +Creates the normalized identifier tables, their graph edges, and nullable links +from existing domain tables. Retained identifier JSON is backfilled on a +best-effort basis and remains available when a legacy value cannot be linked. + +Revision ID: e5f7a9c1b3d2 +Revises: d4e6f8a0b2c4 +Create Date: 2026-07-10 12:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +import uuid +from collections.abc import Sequence # noqa: TC003 +from functools import partial +from typing import TYPE_CHECKING, Any + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.sqlite import CHAR +from sqlalchemy.types import TypeDecorator, Uuid + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy.engine import Dialect + +# revision identifiers, used by Alembic. +revision: str = "e5f7a9c1b3d2" +down_revision: str | None = "d4e6f8a0b2c4" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +class _CustomUUID(TypeDecorator[uuid.UUID]): + """Frozen UUID type matching ``PromptMemoryEntries.id`` across dialects.""" + + impl = CHAR + cache_ok = True + + def load_dialect_impl(self, dialect: Dialect) -> Any: + if dialect.name == "sqlite": + return dialect.type_descriptor(CHAR(36)) + return dialect.type_descriptor(Uuid()) + + def process_bind_param(self, value: Any, dialect: Any) -> str | None: + return str(value) if value is not None else None + + def process_result_value(self, value: Any, dialect: Any) -> uuid.UUID | None: + if value is None: + return None + if isinstance(value, uuid.UUID): + return value + return uuid.UUID(value) + + +def run_best_effort_backfill(*, bind: Any, name: str, backfill: Callable[[], None]) -> None: + """Run a data backfill in a savepoint without blocking the schema upgrade.""" + try: + with bind.begin_nested(): + backfill() + except Exception: + logger.warning(f"{name} backfill failed; leaving new identifier links nullable", exc_info=True) + + +def _run_best_effort_row(*, bind: Any, description: str, operation: Callable[[], None]) -> bool: + """ + Run one backfill row in a savepoint and report whether it succeeded. + + Returns: + bool: Whether the row operation completed successfully. + """ + try: + with bind.begin_nested(): + operation() + except Exception: + logger.warning(description, exc_info=True) + return False + return True + + +def _insert_identifier_link( + *, + bind: Any, + insert: Callable[[dict[str, Any]], str | None], + identifier: dict[str, Any], + update_statement: Any, + update_values: dict[str, Any], +) -> None: + """Insert an identifier graph and update its domain-row link.""" + identifier_hash = insert(identifier) + if not identifier_hash: + return + bind.execute(update_statement, {**update_values, "hash": identifier_hash}) + + +def load_identifier(raw_identifier: Any) -> dict[str, Any] | None: + """ + Load a retained identifier JSON value without importing domain models. + + Returns: + dict[str, Any] | None: The identifier dictionary when it has a usable hash. + """ + try: + value = json.loads(raw_identifier) if isinstance(raw_identifier, str) else raw_identifier + except (TypeError, ValueError): + return None + if not isinstance(value, dict): + return None + identifier_hash = value.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64: + return None + return value + + +def load_identifier_list(raw_identifiers: Any) -> list[dict[str, Any]]: + """ + Load the valid identifiers from a retained JSON list. + + Returns: + list[dict[str, Any]]: Identifier dictionaries carrying usable hashes. + """ + try: + values = json.loads(raw_identifiers) if isinstance(raw_identifiers, str) else raw_identifiers + except (TypeError, ValueError): + return [] + if not isinstance(values, list): + return [] + return [identifier for value in values if (identifier := load_identifier(value)) is not None] + + +class IdentifierGraphInserter: + """Best-effort inserter for the frozen flat identifier JSON shape.""" + + _TABLES = ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ConverterIdentifiers", + "ScenarioIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ) + + def __init__(self, *, bind: Any) -> None: + """Initialize the inserter from tables available at this migration revision.""" + self._bind = bind + table_names = set(sa.inspect(bind).get_table_names()) + self._hashes = { + table: set(bind.execute(sa.text(f'SELECT hash FROM "{table}"')).scalars()) + for table in self._TABLES + if table in table_names + } + + def insert_target(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a target graph. + + Returns: + str | None: The stored hash when successful. + """ + children = self._children(identifier, "targets") + child_hashes = [child_hash for child in children if (child_hash := self.insert_target(child))] + identifier_hash = self._insert_identifier( + table="TargetIdentifiers", + identifier=identifier, + promoted=( + "endpoint", + "model_name", + "underlying_model_name", + "temperature", + "top_p", + "max_requests_per_minute", + "supported_auth_modes", + ), + ) + if identifier_hash: + self._insert_edges( + table="TargetIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_scorer(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scorer graph. + + Returns: + str | None: The stored hash when successful. + """ + prompt_target = self._child(identifier, "prompt_target", aliases=("chat_target",)) + prompt_target_hash = self.insert_target(prompt_target) if prompt_target else None + sub_scorers = self._children(identifier, "sub_scorers", aliases=("scorers",)) + child_hashes = [child_hash for child in sub_scorers if (child_hash := self.insert_scorer(child))] + identifier_hash = self._insert_identifier( + table="ScorerIdentifiers", + identifier=identifier, + promoted=("scorer_type", "score_aggregator"), + extra={"prompt_target_hash": prompt_target_hash}, + ) + if identifier_hash: + self._insert_edges( + table="ScorerIdentifierChildren", + parent_column="parent_hash", + parent_hash=identifier_hash, + child_column="child_hash", + child_hashes=child_hashes, + ) + return identifier_hash + + def insert_converter(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a converter graph. + + Returns: + str | None: The stored hash when successful. + """ + converter_target = self._child(identifier, "converter_target") + sub_converter = self._child(identifier, "sub_converter") + return self._insert_identifier( + table="ConverterIdentifiers", + identifier=identifier, + promoted=("supported_input_types", "supported_output_types"), + extra={ + "converter_target_hash": self.insert_target(converter_target) if converter_target else None, + "sub_converter_hash": self.insert_converter(sub_converter) if sub_converter else None, + }, + ) + + def insert_scenario(self, identifier: dict[str, Any]) -> str | None: + """ + Insert a scenario graph. + + Returns: + str | None: The stored hash when successful. + """ + objective_target = self._child(identifier, "objective_target") + objective_scorer = self._child(identifier, "objective_scorer") + return self._insert_identifier( + table="ScenarioIdentifiers", + identifier=identifier, + promoted=("version", "techniques", "datasets"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + + def insert_atomic_attack(self, identifier: dict[str, Any]) -> str | None: + """ + Insert an atomic attack graph. + + Returns: + str | None: The stored hash when successful. + """ + attack_technique = self._child(identifier, "attack_technique") + seeds = self._children(identifier, "seed_identifiers") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AtomicAttackIdentifiers", + identifier=identifier, + extra={ + "attack_technique_identifier_hash": ( + self._insert_attack_technique(attack_technique) if attack_technique else None + ) + }, + ) + if identifier_hash: + self._insert_edges( + table="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack_technique(self, identifier: dict[str, Any]) -> str | None: + attack = self._child(identifier, "attack") + seeds = self._children(identifier, "technique_seeds") + seed_hashes = [seed_hash for seed in seeds if (seed_hash := self._insert_seed(seed))] + identifier_hash = self._insert_identifier( + table="AttackTechniqueIdentifiers", + identifier=identifier, + extra={"attack_identifier_hash": self._insert_attack(attack) if attack else None}, + ) + if identifier_hash: + self._insert_edges( + table="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_hash=identifier_hash, + child_column="seed_identifier_hash", + child_hashes=seed_hashes, + ) + return identifier_hash + + def _insert_attack(self, identifier: dict[str, Any]) -> str | None: + objective_target = self._child(identifier, "objective_target") + adversarial_chat = self._child(identifier, "adversarial_chat") + objective_scorer = self._child(identifier, "objective_scorer") + request_hashes = [ + value for item in self._children(identifier, "request_converters") if (value := self.insert_converter(item)) + ] + response_hashes = [ + value + for item in self._children(identifier, "response_converters") + if (value := self.insert_converter(item)) + ] + identifier_hash = self._insert_identifier( + table="AttackIdentifiers", + identifier=identifier, + promoted=("adversarial_system_prompt", "adversarial_seed_prompt"), + extra={ + "objective_target_hash": self.insert_target(objective_target) if objective_target else None, + "adversarial_chat_hash": self.insert_target(adversarial_chat) if adversarial_chat else None, + "objective_scorer_hash": self.insert_scorer(objective_scorer) if objective_scorer else None, + }, + ) + if identifier_hash: + self._insert_edges( + table="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=request_hashes, + ) + self._insert_edges( + table="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_hash=identifier_hash, + child_column="converter_identifier_hash", + child_hashes=response_hashes, + ) + return identifier_hash + + def _insert_seed(self, identifier: dict[str, Any]) -> str | None: + return self._insert_identifier( + table="SeedIdentifiers", + identifier=identifier, + promoted=("value", "value_sha256", "data_type", "dataset_name", "is_general_technique"), + ) + + def _insert_identifier( + self, + *, + table: str, + identifier: dict[str, Any], + promoted: Sequence[str] = (), + extra: dict[str, Any] | None = None, + ) -> str | None: + identifier_hash = identifier.get("hash") + if not isinstance(identifier_hash, str) or len(identifier_hash) != 64 or table not in self._hashes: + return None + if identifier_hash in self._hashes[table]: + return identifier_hash + values: dict[str, Any] = { + "hash": identifier_hash, + "class_name": identifier.get("class_name"), + "class_module": identifier.get("class_module"), + "identifier_json": json.dumps(identifier, sort_keys=True), + "pyrit_version": identifier.get("pyrit_version"), + } + values.update({name: self._json_value(identifier.get(name)) for name in promoted}) + values.update(extra or {}) + columns = list(values) + statement = sa.text( + f'INSERT INTO "{table}" ({", ".join(columns)}) VALUES ({", ".join(f":{column}" for column in columns)})' + ) + self._bind.execute(statement, values) + self._hashes[table].add(identifier_hash) + return identifier_hash + + def _insert_edges( + self, + *, + table: str, + parent_column: str, + parent_hash: str, + child_column: str, + child_hashes: Sequence[str], + ) -> None: + select_statement = sa.text( + f'SELECT "{child_column}" FROM "{table}" WHERE "{parent_column}" = :parent_hash AND position = :position' + ) + statement = sa.text( + f'INSERT INTO "{table}" ("{parent_column}", position, "{child_column}") ' + f"VALUES (:parent_hash, :position, :child_hash)" + ) + for position, child_hash in enumerate(child_hashes): + parameters = {"parent_hash": parent_hash, "position": position} + existing_child_hash = self._bind.execute(select_statement, parameters).scalar_one_or_none() + if existing_child_hash == child_hash: + continue + if existing_child_hash is not None: + raise ValueError( + f"Conflicting {table} edge for parent {parent_hash!r} at position {position}: " + f"stored child {existing_child_hash!r}, retained child {child_hash!r}." + ) + self._bind.execute( + statement, + {**parameters, "child_hash": child_hash}, + ) + + @staticmethod + def _child( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> dict[str, Any] | None: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, dict): + return load_identifier(value) + return None + + @staticmethod + def _children( + identifier: dict[str, Any], + name: str, + aliases: Sequence[str] = (), + ) -> list[dict[str, Any]]: + children = identifier.get("children") + children = children if isinstance(children, dict) else {} + for key in (name, *aliases): + value = identifier.get(key, children.get(key)) + if isinstance(value, list): + return [child for item in value if (child := load_identifier(item)) is not None] + return [] + + @staticmethod + def _json_value(value: Any) -> Any: + return json.dumps(value) if isinstance(value, (list, dict)) else value + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "TargetIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), + sa.Column("endpoint", sa.String(), nullable=True), + sa.Column("model_name", sa.String(), nullable=True), + sa.Column("underlying_model_name", sa.String(), nullable=True), + sa.Column("temperature", sa.Float(), nullable=True), + sa.Column("top_p", sa.Float(), nullable=True), + sa.Column("max_requests_per_minute", sa.Integer(), nullable=True), + sa.Column("supported_auth_modes", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + # Self-referential pivot mapping a multi-target to its inner target identifiers. + # Both endpoints are content hashes into TargetIdentifiers; ``position`` preserves + # the parent's ``targets`` list order. Named FK constraints for SQL Server / batch + # portability. + op.create_table( + "TargetIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_child_hash" + ), + ) + + op.create_table( + "ScorerIdentifiers", + *_common_columns(), + sa.Column("scorer_type", sa.String(), nullable=True), + sa.Column("score_aggregator", sa.String(), nullable=True), + sa.Column("prompt_target_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["prompt_target_hash"], ["TargetIdentifiers.hash"], name="fk_scorer_identifiers_prompt_target_hash" + ), + ) + op.create_table( + "ScorerIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["ScorerIdentifiers.hash"], name="fk_scorer_identifier_children_child_hash" + ), + ) + op.create_table( + "ScenarioIdentifiers", + *_common_columns(), + sa.Column("version", sa.Integer(), nullable=True), + sa.Column("techniques", sa.JSON(), nullable=True), + sa.Column("datasets", sa.JSON(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["objective_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_scenario_identifiers_objective_target_hash", + ), + sa.ForeignKeyConstraint( + ["objective_scorer_hash"], + ["ScorerIdentifiers.hash"], + name="fk_scenario_identifiers_objective_scorer_hash", + ), + ) + op.create_table( + "ConverterIdentifiers", + *_common_columns(), + sa.Column("supported_input_types", sa.JSON(), nullable=True), + sa.Column("supported_output_types", sa.JSON(), nullable=True), + sa.Column("converter_target_hash", sa.String(64), nullable=True), + sa.Column("sub_converter_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint( + ["converter_target_hash"], + ["TargetIdentifiers.hash"], + name="fk_converter_identifiers_converter_target_hash", + ), + sa.ForeignKeyConstraint( + ["sub_converter_hash"], + ["ConverterIdentifiers.hash"], + name="fk_converter_identifiers_sub_converter_hash", + ), + ) + op.create_table( + "PromptConverterIdentifiers", + sa.Column("prompt_memory_entry_id", _CustomUUID(), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("converter_identifier_hash", sa.String(64), nullable=False), + sa.ForeignKeyConstraint( + ["prompt_memory_entry_id"], + ["PromptMemoryEntries.id"], + name="fk_prompt_converter_identifiers_prompt_memory_entry_id", + ), + sa.ForeignKeyConstraint( + ["converter_identifier_hash"], + ["ConverterIdentifiers.hash"], + name="fk_prompt_converter_identifiers_converter_identifier_hash", + ), + sa.PrimaryKeyConstraint("prompt_memory_entry_id", "position"), + ) + _create_attack_identifier_tables() + + # Batch op for SQLite portability (no ALTER TABLE ADD FOREIGN KEY on SQLite). + # The FK constraint must be named explicitly: Alembic batch mode rejects an + # unnamed constraint. + with op.batch_alter_table("Conversations") as batch_op: + batch_op.add_column(sa.Column("target_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_conversations_target_identifier_hash", + "TargetIdentifiers", + ["target_identifier_hash"], + ["hash"], + ) + + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.add_column(sa.Column("scorer_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_score_entries_scorer_identifier_hash", + "ScorerIdentifiers", + ["scorer_identifier_hash"], + ["hash"], + ) + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.add_column(sa.Column("scenario_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_scenario_result_entries_scenario_identifier_hash", + "ScenarioIdentifiers", + ["scenario_identifier_hash"], + ["hash"], + ) + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.add_column(sa.Column("atomic_attack_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_attack_result_entries_atomic_attack_identifier_hash", + "AtomicAttackIdentifiers", + ["atomic_attack_identifier_hash"], + ["hash"], + ) + + bind = op.get_bind() + for name, backfill in ( + ("TargetIdentifiers", _backfill_target_identifiers), + ("ScorerIdentifiers", _backfill_scorer_identifiers), + ("ScenarioIdentifiers", _backfill_scenario_identifiers), + ("ConverterIdentifiers", _backfill_converter_identifiers), + ("AttackIdentifiers", _backfill_attack_identifiers), + ): + run_best_effort_backfill(bind=bind, name=name, backfill=backfill) + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.drop_constraint("fk_attack_result_entries_atomic_attack_identifier_hash", type_="foreignkey") + batch_op.drop_column("atomic_attack_identifier_hash") + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.drop_constraint("fk_scenario_result_entries_scenario_identifier_hash", type_="foreignkey") + batch_op.drop_column("scenario_identifier_hash") + with op.batch_alter_table("ScoreEntries") as batch_op: + batch_op.drop_constraint("fk_score_entries_scorer_identifier_hash", type_="foreignkey") + batch_op.drop_column("scorer_identifier_hash") + with op.batch_alter_table("Conversations") as batch_op: + batch_op.drop_constraint("fk_conversations_target_identifier_hash", type_="foreignkey") + batch_op.drop_column("target_identifier_hash") + + op.drop_table("AtomicAttackSeedIdentifiers") + op.drop_table("AtomicAttackIdentifiers") + op.drop_table("AttackTechniqueSeedIdentifiers") + op.drop_table("AttackTechniqueIdentifiers") + op.drop_table("AttackResponseConverterIdentifiers") + op.drop_table("AttackRequestConverterIdentifiers") + op.drop_table("AttackIdentifiers") + op.drop_table("SeedIdentifiers") + op.drop_table("PromptConverterIdentifiers") + op.drop_table("ConverterIdentifiers") + op.drop_table("ScenarioIdentifiers") + op.drop_table("ScorerIdentifierChildren") + op.drop_table("ScorerIdentifiers") + op.drop_table("TargetIdentifierChildren") + op.drop_table("TargetIdentifiers") + + +def _common_columns() -> tuple[sa.Column[Any], ...]: + return ( + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=True), + sa.Column("class_module", sa.String(), nullable=True), + sa.Column("identifier_json", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + +def _create_ordered_edge_table( + *, + table_name: str, + parent_column: str, + parent_table: str, + child_column: str, + child_table: str, +) -> None: + op.create_table( + table_name, + sa.Column(parent_column, sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column(child_column, sa.String(64), nullable=False), + sa.ForeignKeyConstraint([parent_column], [f"{parent_table}.hash"]), + sa.ForeignKeyConstraint([child_column], [f"{child_table}.hash"]), + sa.PrimaryKeyConstraint(parent_column, "position"), + ) + + +def _create_attack_identifier_tables() -> None: + op.create_table( + "SeedIdentifiers", + *_common_columns(), + sa.Column("value", sa.Unicode(), nullable=True), + sa.Column("value_sha256", sa.String(), nullable=True), + sa.Column("data_type", sa.String(), nullable=True), + sa.Column("dataset_name", sa.String(), nullable=True), + sa.Column("is_general_technique", sa.Boolean(), nullable=True), + ) + op.create_table( + "AttackIdentifiers", + *_common_columns(), + sa.Column("adversarial_system_prompt", sa.Unicode(), nullable=True), + sa.Column("adversarial_seed_prompt", sa.Unicode(), nullable=True), + sa.Column("objective_target_hash", sa.String(64), nullable=True), + sa.Column("adversarial_chat_hash", sa.String(64), nullable=True), + sa.Column("objective_scorer_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["objective_target_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["adversarial_chat_hash"], ["TargetIdentifiers.hash"]), + sa.ForeignKeyConstraint(["objective_scorer_hash"], ["ScorerIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackRequestConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + _create_ordered_edge_table( + table_name="AttackResponseConverterIdentifiers", + parent_column="attack_identifier_hash", + parent_table="AttackIdentifiers", + child_column="converter_identifier_hash", + child_table="ConverterIdentifiers", + ) + op.create_table( + "AttackTechniqueIdentifiers", + *_common_columns(), + sa.Column("attack_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["attack_identifier_hash"], ["AttackIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AttackTechniqueSeedIdentifiers", + parent_column="attack_technique_identifier_hash", + parent_table="AttackTechniqueIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + op.create_table( + "AtomicAttackIdentifiers", + *_common_columns(), + sa.Column("attack_technique_identifier_hash", sa.String(64), nullable=True), + sa.ForeignKeyConstraint(["attack_technique_identifier_hash"], ["AttackTechniqueIdentifiers.hash"]), + ) + _create_ordered_edge_table( + table_name="AtomicAttackSeedIdentifiers", + parent_column="atomic_attack_identifier_hash", + parent_table="AtomicAttackIdentifiers", + child_column="seed_identifier_hash", + child_table="SeedIdentifiers", + ) + + +def _insert_converter_links( + *, + bind: Any, + inserter: IdentifierGraphInserter, + link_statement: Any, + prompt_id: Any, + stored_identifiers: Any, + pyrit_version: str | None, +) -> None: + """Insert converter graphs and their ordered links for one prompt row.""" + for position, identifier in enumerate(load_identifier_list(stored_identifiers)): + if identifier.get("pyrit_version") is None: + identifier = {**identifier, "pyrit_version": pyrit_version} + identifier_hash = inserter.insert_converter(identifier) + if identifier_hash: + bind.execute( + link_statement, + { + "prompt_memory_entry_id": prompt_id, + "position": position, + "converter_identifier_hash": identifier_hash, + }, + ) + + +def _backfill_target_identifiers() -> None: + """ + Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set + ``Conversations.target_identifier_hash``. + + For every ``Conversations`` row with a non-null ``target_identifier`` JSON, + load the retained ``TargetIdentifier`` shape and its stored hash, insert the + deduped ``TargetIdentifiers`` row if absent -- recursing into any inner + ``targets`` first so the child edge foreign keys resolve -- record the + ``parent_hash -> child_hash`` edges, and point the conversation's + ``target_identifier_hash`` at the top-level row. Idempotent: hashes already present + are not re-inserted. Rows whose stored target cannot be reconstructed are logged and + skipped rather than aborting the upgrade. + """ + bind = op.get_bind() + rows = bind.execute( + sa.text( + 'SELECT conversation_id, target_identifier FROM "Conversations" ' + "WHERE target_identifier IS NOT NULL ORDER BY conversation_id" + ) + ).fetchall() + + update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') + inserter = IdentifierGraphInserter(bind=bind) + linked = 0 + skipped = 0 + for conversation_id, raw_target in rows: + identifier = load_identifier(raw_target) + if identifier is None: + skipped += 1 + continue + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_target, + identifier=identifier, + update_statement=update_stmt, + update_values={"cid": conversation_id}, + ) + if _run_best_effort_row( + bind=bind, + description=f"TargetIdentifiers backfill skipped conversation {conversation_id!r}", + operation=operation, + ): + linked += 1 + else: + skipped += 1 + inserter = IdentifierGraphInserter(bind=bind) + + if linked or skipped: + logger.info(f"TargetIdentifiers backfill linked {linked} conversation(s); skipped {skipped}.") + + +def _backfill_scorer_identifiers() -> None: + """Backfill scorer rows and score foreign keys from retained JSON.""" + bind = op.get_bind() + score_rows = bind.execute( + sa.text( + 'SELECT id, scorer_class_identifier FROM "ScoreEntries" ' + "WHERE scorer_class_identifier IS NOT NULL ORDER BY id" + ) + ).fetchall() + score_update = sa.text('UPDATE "ScoreEntries" SET scorer_identifier_hash = :hash WHERE id = :id') + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for score_id, raw_scorer in score_rows: + identifier = load_identifier(raw_scorer) + if identifier is None: + skipped += 1 + continue + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_scorer, + identifier=identifier, + update_statement=score_update, + update_values={"id": score_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ScorerIdentifiers backfill: could not reconstruct scorer for score {score_id}", + operation=operation, + ): + skipped += 1 + inserter = IdentifierGraphInserter(bind=bind) + if skipped: + logger.warning(f"ScorerIdentifiers backfill skipped {skipped} score row(s)") + + +def _backfill_scenario_identifiers() -> None: + """Backfill scenario rows and result foreign keys from retained JSON.""" + bind = op.get_bind() + result_rows = bind.execute( + sa.text( + 'SELECT id, scenario_identifier FROM "ScenarioResultEntries" ' + "WHERE scenario_identifier IS NOT NULL ORDER BY id" + ) + ).fetchall() + update_stmt = sa.text('UPDATE "ScenarioResultEntries" SET scenario_identifier_hash = :hash WHERE id = :id') + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for result_id, raw_scenario in result_rows: + identifier = load_identifier(raw_scenario) + if identifier is None: + skipped += 1 + continue + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_scenario, + identifier=identifier, + update_statement=update_stmt, + update_values={"id": result_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ScenarioIdentifiers backfill: could not reconstruct scenario for result {result_id}", + operation=operation, + ): + skipped += 1 + inserter = IdentifierGraphInserter(bind=bind) + if skipped: + logger.warning(f"ScenarioIdentifiers backfill skipped {skipped} scenario result row(s)") + + +def _backfill_converter_identifiers() -> None: + """Materialize converter graphs and prompt associations from retained JSON.""" + bind = op.get_bind() + prompt_rows = bind.execute( + sa.text( + 'SELECT id, converter_identifiers, pyrit_version FROM "PromptMemoryEntries" ' + "WHERE converter_identifiers IS NOT NULL ORDER BY id" + ) + ).fetchall() + link_insert = sa.text( + 'INSERT INTO "PromptConverterIdentifiers" ' + "(prompt_memory_entry_id, position, converter_identifier_hash) " + "VALUES (:prompt_memory_entry_id, :position, :converter_identifier_hash)" + ) + inserter = IdentifierGraphInserter(bind=bind) + skipped = 0 + for prompt_id, stored_identifiers, pyrit_version in prompt_rows: + operation = partial( + _insert_converter_links, + bind=bind, + inserter=inserter, + link_statement=link_insert, + prompt_id=prompt_id, + stored_identifiers=stored_identifiers, + pyrit_version=pyrit_version, + ) + if not _run_best_effort_row( + bind=bind, + description=f"ConverterIdentifiers backfill: could not reconstruct converters for prompt {prompt_id}", + operation=operation, + ): + skipped += 1 + inserter = IdentifierGraphInserter(bind=bind) + if skipped: + logger.warning(f"ConverterIdentifiers backfill skipped {skipped} prompt row(s)") + + +def _backfill_attack_identifiers() -> None: + """Backfill attack identifier graphs and result links from retained JSON.""" + bind = op.get_bind() + result_rows = bind.execute( + sa.text( + 'SELECT id, atomic_attack_identifier FROM "AttackResultEntries" ' + "WHERE atomic_attack_identifier IS NOT NULL ORDER BY id" + ) + ).fetchall() + inserter = IdentifierGraphInserter(bind=bind) + update_stmt = sa.text('UPDATE "AttackResultEntries" SET atomic_attack_identifier_hash = :hash WHERE id = :id') + skipped = 0 + for result_id, raw_identifier in result_rows: + identifier = load_identifier(raw_identifier) + if identifier is None: + skipped += 1 + continue + operation = partial( + _insert_identifier_link, + bind=bind, + insert=inserter.insert_atomic_attack, + identifier=identifier, + update_statement=update_stmt, + update_values={"id": result_id}, + ) + if not _run_best_effort_row( + bind=bind, + description=f"Attack identifier backfill could not reconstruct result {result_id}", + operation=operation, + ): + skipped += 1 + inserter = IdentifierGraphInserter(bind=bind) + if skipped: + logger.warning(f"Attack identifier backfill skipped {skipped} attack result row(s)") 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..98f96a75f1 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, ) @@ -61,6 +80,7 @@ Model = TypeVar("Model") +IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) class MemoryInterface(abc.ABC): @@ -395,20 +415,39 @@ def add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece] self._validate_persistable_conversation_ids(message_pieces=pieces_to_insert) self._add_message_pieces_to_memory(message_pieces=pieces_to_insert) - @abc.abstractmethod def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece]) -> None: """ Persist already-validated message pieces to the backing store. Called by ``add_message_pieces_to_memory`` after ``not_in_memory`` pieces are - filtered out and conversation_ids are validated. Implementations only translate - the pieces into storage rows and insert them; they must not re-filter or - re-validate. + filtered out and conversation_ids are validated. Args: message_pieces (Sequence[MessagePiece]): Persistable pieces (none flagged ``not_in_memory``), each carrying a non-empty ``conversation_id``. + + Raises: + SQLAlchemyError: If the message pieces or converter identifiers cannot be persisted. """ + entries = [PromptMemoryEntry(entry=piece) for piece in message_pieces] + with closing(self.get_session()) as session: + 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 +498,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 +522,427 @@ 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. + + Dependencies are persisted before the target row, whose ordered child edges + reference those rows by content hash. Identifier rows are immutable and keyed + by their content hash, so an identical target reused across many conversations + maps to a single row. + + If the row already exists it was fully persisted before (children and edges + included, since rows are immutable), so this returns early. Otherwise the row and + its child edges are inserted inside a savepoint. If an ``IntegrityError`` occurs, + it is treated as a concurrent duplicate only when a fresh lookup confirms that + the identifier hash now exists; all other integrity failures are re-raised. + + Args: + session (Any): The active SQLAlchemy session (the caller's transaction). + 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.add(entry_type.from_domain_model(identifier)) + session.flush() + except IntegrityError: + with session.no_autoflush: + existing_entry = session.get(entry_type, identifier.hash, populate_existing=True) + if existing_entry is None: + raise + + @staticmethod + def _get_identifier_entry_type(identifier: ComponentIdentifier) -> type[ComponentIdentifierEntry[Any]]: + 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)) + + def _get_identifiers( + self, + *, + identifier_type: type[IdentifierModel], + entry_type: type[ComponentIdentifierEntry[Any]], + identifier_hashes: Sequence[str] | None, + filters: dict[str, Any], + ) -> Sequence[IdentifierModel]: + if identifier_hashes is not None and not identifier_hashes: + return [] + + list_filters = {name: value for name, value in filters.items() if isinstance(value, list)} + conditions = [ + getattr(entry_type, name) == value + for name, value in filters.items() + if value is not None and name not in list_filters + ] + if identifier_hashes is not None: + entries = self._execute_batched_query( + entry_type, + batch_column=entry_type.hash, + batch_values=identifier_hashes, + other_conditions=conditions, + order_by=entry_type.hash, + ) + else: + entries = self._query_entries( + entry_type, + conditions=and_(*conditions) if conditions else None, + order_by=entry_type.hash, + ) + + entries = [ + entry + for entry in entries + if all( + getattr(entry, name) is not None and sorted(getattr(entry, name)) == sorted(value) + for name, value in list_filters.items() + ) + ] + identifiers: list[IdentifierModel] = [] + seen_hashes: set[str] = set() + for entry in sorted(entries, key=lambda item: item.hash): + if entry.hash in seen_hashes: + continue + if entry.identifier_json is None: + raise ValueError(f"Identifier row {entry.hash} in {entry_type.__tablename__} has no identifier JSON.") + identifier = identifier_type.model_validate(entry.identifier_json) + if identifier.hash != entry.hash: + raise ValueError( + f"Identifier row {entry.hash} in {entry_type.__tablename__} does not match its stored JSON hash." + ) + identifiers.append(identifier) + seen_hashes.add(entry.hash) + return identifiers + + def get_target_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + endpoint: str | None = None, + model_name: str | None = None, + underlying_model_name: str | None = None, + temperature: float | None = None, + top_p: float | None = None, + max_requests_per_minute: int | None = None, + supported_auth_modes: Sequence[str] | None = None, + ) -> Sequence[TargetIdentifier]: + """ + Retrieve target identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + endpoint (str | None): Target endpoint to match. + model_name (str | None): Target model name to match. + underlying_model_name (str | None): Underlying model name to match. + temperature (float | None): Temperature to match. + top_p (float | None): Top-p value to match. + max_requests_per_minute (int | None): Request limit to match. + supported_auth_modes (Sequence[str] | None): Authentication modes to match exactly, in any order. + + Returns: + Sequence[TargetIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=TargetIdentifier, + entry_type=TargetIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "endpoint": endpoint, + "model_name": model_name, + "underlying_model_name": underlying_model_name, + "temperature": temperature, + "top_p": top_p, + "max_requests_per_minute": max_requests_per_minute, + "supported_auth_modes": list(supported_auth_modes) if supported_auth_modes is not None else None, + }, + ) + + def get_converter_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + supported_input_types: Sequence[str] | None = None, + supported_output_types: Sequence[str] | None = None, + converter_target_hash: str | None = None, + sub_converter_hash: str | None = None, + ) -> Sequence[ConverterIdentifier]: + """ + Retrieve converter identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + supported_input_types (Sequence[str] | None): Input types to match exactly, in any order. + supported_output_types (Sequence[str] | None): Output types to match exactly, in any order. + converter_target_hash (str | None): Converter target hash to match. + sub_converter_hash (str | None): Nested converter hash to match. + + Returns: + Sequence[ConverterIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ConverterIdentifier, + entry_type=ConverterIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "supported_input_types": (list(supported_input_types) if supported_input_types is not None else None), + "supported_output_types": ( + list(supported_output_types) if supported_output_types is not None else None + ), + "converter_target_hash": converter_target_hash, + "sub_converter_hash": sub_converter_hash, + }, + ) + + def get_scorer_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + scorer_type: str | None = None, + score_aggregator: str | None = None, + prompt_target_hash: str | None = None, + ) -> Sequence[ScorerIdentifier]: + """ + Retrieve scorer identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + scorer_type (str | None): Scorer type to match. + score_aggregator (str | None): Score aggregator to match. + prompt_target_hash (str | None): Scorer target hash to match. + + Returns: + Sequence[ScorerIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ScorerIdentifier, + entry_type=ScorerIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "scorer_type": scorer_type, + "score_aggregator": score_aggregator, + "prompt_target_hash": prompt_target_hash, + }, + ) + + def get_scenario_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + version: int | None = None, + techniques: Sequence[str] | None = None, + datasets: Sequence[str] | None = None, + objective_target_hash: str | None = None, + objective_scorer_hash: str | None = None, + ) -> Sequence[ScenarioIdentifier]: + """ + Retrieve scenario identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + version (int | None): Scenario definition version to match. + techniques (Sequence[str] | None): Technique names to match exactly, in any order. + datasets (Sequence[str] | None): Dataset names to match exactly, in any order. + objective_target_hash (str | None): Objective target hash to match. + objective_scorer_hash (str | None): Objective scorer hash to match. + + Returns: + Sequence[ScenarioIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=ScenarioIdentifier, + entry_type=ScenarioIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "version": version, + "techniques": list(techniques) if techniques is not None else None, + "datasets": list(datasets) if datasets is not None else None, + "objective_target_hash": objective_target_hash, + "objective_scorer_hash": objective_scorer_hash, + }, + ) + + def get_seed_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + value: str | None = None, + value_sha256: str | None = None, + data_type: str | None = None, + dataset_name: str | None = None, + is_general_technique: bool | None = None, + ) -> Sequence[SeedIdentifier]: + """ + Retrieve seed identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + value (str | None): Seed value to match. + value_sha256 (str | None): Seed value hash to match. + data_type (str | None): Seed data type to match. + dataset_name (str | None): Dataset name to match. + is_general_technique (bool | None): General-technique flag to match. + + Returns: + Sequence[SeedIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=SeedIdentifier, + entry_type=SeedIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "value": value, + "value_sha256": value_sha256, + "data_type": data_type, + "dataset_name": dataset_name, + "is_general_technique": is_general_technique, + }, + ) + + def get_attack_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + adversarial_system_prompt: str | None = None, + adversarial_seed_prompt: str | None = None, + objective_target_hash: str | None = None, + adversarial_chat_hash: str | None = None, + objective_scorer_hash: str | None = None, + ) -> Sequence[AttackIdentifier]: + """ + Retrieve attack identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + adversarial_system_prompt (str | None): Adversarial system prompt to match. + adversarial_seed_prompt (str | None): Adversarial seed prompt to match. + objective_target_hash (str | None): Objective target hash to match. + adversarial_chat_hash (str | None): Adversarial chat target hash to match. + objective_scorer_hash (str | None): Objective scorer hash to match. + + Returns: + Sequence[AttackIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AttackIdentifier, + entry_type=AttackIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "adversarial_system_prompt": adversarial_system_prompt, + "adversarial_seed_prompt": adversarial_seed_prompt, + "objective_target_hash": objective_target_hash, + "adversarial_chat_hash": adversarial_chat_hash, + "objective_scorer_hash": objective_scorer_hash, + }, + ) + + def get_attack_technique_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + attack_identifier_hash: str | None = None, + ) -> Sequence[AttackTechniqueIdentifier]: + """ + Retrieve attack technique identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + attack_identifier_hash (str | None): Attack identifier hash to match. + + Returns: + Sequence[AttackTechniqueIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AttackTechniqueIdentifier, + entry_type=AttackTechniqueIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "attack_identifier_hash": attack_identifier_hash, + }, + ) + + def get_atomic_attack_identifiers( + self, + *, + identifier_hashes: Sequence[str] | None = None, + class_name: str | None = None, + attack_technique_identifier_hash: str | None = None, + ) -> Sequence[AtomicAttackIdentifier]: + """ + Retrieve atomic attack identifiers using exact normalized-column filters. + + Args: + identifier_hashes (Sequence[str] | None): Content hashes to include. + class_name (str | None): Component class name to match. + attack_technique_identifier_hash (str | None): Attack technique hash to match. + + Returns: + Sequence[AtomicAttackIdentifier]: Matching identifiers ordered by content hash. + """ + return self._get_identifiers( + identifier_type=AtomicAttackIdentifier, + entry_type=AtomicAttackIdentifierEntry, + identifier_hashes=identifier_hashes, + filters={ + "class_name": class_name, + "attack_technique_identifier_hash": attack_technique_identifier_hash, + }, + ) + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ @@ -804,6 +1271,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 +1285,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 +2260,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 +2624,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..59a93d5a4f 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, get_args, get_origin 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,626 @@ 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_hash_attr: str + 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 projection columns the way ``TargetIdentifier`` inherits + ``ComponentIdentifier``. The content ``hash`` is the natural, dedupable primary + key. Runtime writes populate the descriptive fields and full ``identifier_json``; + they remain nullable so best-effort migration backfills can preserve partial + legacy identifiers. Rows are immutable — the same content always maps to the + same hash, so a given identifier reused across rows is stored once. + + Subclasses declare their promoted query columns and implement the + ``DomainBackedEntry.from_domain_model`` seam to map their strongly-typed identifier + 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 | None] = mapped_column(String, nullable=True) + class_module: Mapped[str | None] = mapped_column(String, nullable=True) + #: Full flat ``model_dump()`` of the identifier. Source of truth on reload. + identifier_json: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) + #: Version that first wrote this content-addressed row. Nullable for backwards + #: compatibility with existing databases. + pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) + + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Validate that concrete entries persist every promoted scalar field. + + Raises: + TypeError: If the entry omits its identifier type or a mapped scalar column. + """ + super().__init_subclass__(**kwargs) + if cls.__dict__.get("__abstract__", False): + return + + identifier_type = next( + ( + get_args(base)[0] + for base in getattr(cls, "__orig_bases__", ()) + if get_origin(base) is ComponentIdentifierEntry + ), + None, + ) + if not isinstance(identifier_type, type) or not issubclass(identifier_type, ComponentIdentifier): + raise TypeError(f"{cls.__name__} must declare its ComponentIdentifier domain model type.") + + missing_columns = set(identifier_type.promoted_scalar_field_names()) - set(cls.__table__.columns.keys()) + if missing_columns: + names = ", ".join(sorted(missing_columns)) + raise TypeError(f"{cls.__name__} has no mapped column for promoted scalar field(s): {names}.") + + @classmethod + def from_domain_model(cls, domain_model: T) -> Self: + """ + 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. + """ + 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) + cls._attach_child_relationship_rows(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) -> 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() + setattr(edge_row, spec.edge_child_hash_attr, child_identifier.hash) + if spec.edge_position_attr: + setattr(edge_row, spec.edge_position_attr, position) + edge_rows.append(edge_row) + + +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_hash_attr="child_hash", + 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", + cascade="all, delete-orphan", + ) + + +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 with its child hash by ``TargetIdentifierEntry.from_domain_model`` + after ``MemoryInterface._persist_target_identifier`` has persisted the child row. + It is a plain ``Base`` row rather than a ``DomainBackedEntry`` because an edge has + no standalone domain model. + """ + + __tablename__ = "TargetIdentifierChildren" + __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_hash_attr="child_hash", + 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", + cascade="all, delete-orphan", + ) + + +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 +1000,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 +1011,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 +1030,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 +1091,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 +1125,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 +1510,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 +1566,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 +1583,15 @@ 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 ) - self.atomic_attack_identifier = ( - entry.atomic_attack_identifier.model_dump() if entry.atomic_attack_identifier else None - ) + entry.atomic_attack_identifier = atomic_attack_identifier + self.atomic_attack_identifier = atomic_attack_identifier.model_dump() if atomic_attack_identifier else None + self.atomic_attack_identifier_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 +1796,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 +1845,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_identifiers.py b/tests/unit/memory/memory_interface/test_interface_identifiers.py new file mode 100644 index 0000000000..8f638655e7 --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_identifiers.py @@ -0,0 +1,275 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from contextlib import closing +from dataclasses import dataclass + +import pytest + +from pyrit.memory import MemoryInterface +from pyrit.memory.memory_models import TargetIdentifierEntry +from pyrit.models import ( + AtomicAttackIdentifier, + AttackIdentifier, + AttackTechniqueIdentifier, + Conversation, + ConverterIdentifier, + ScenarioIdentifier, + ScorerIdentifier, + SeedIdentifier, + TargetIdentifier, +) + + +@dataclass(frozen=True) +class IdentifierGraph: + objective_target: TargetIdentifier + adversarial_target: TargetIdentifier + nested_converter: ConverterIdentifier + converter: ConverterIdentifier + scorer: ScorerIdentifier + scenario: ScenarioIdentifier + technique_seed: SeedIdentifier + dataset_seed: SeedIdentifier + attack: AttackIdentifier + technique: AttackTechniqueIdentifier + atomic_attack: AtomicAttackIdentifier + + +@pytest.fixture +def identifier_graph(sqlite_instance: MemoryInterface) -> IdentifierGraph: + objective_target = TargetIdentifier( + class_name="ObjectiveTarget", + class_module="tests.unit.memory", + endpoint="https://objective.test", + model_name="objective-model", + temperature=0.5, + ) + adversarial_target = TargetIdentifier( + class_name="AdversarialTarget", + class_module="tests.unit.memory", + endpoint="https://adversarial.test", + model_name="adversarial-model", + ) + nested_converter = ConverterIdentifier( + class_name="NestedConverter", + class_module="tests.unit.memory", + supported_input_types=["text", "image_path"], + supported_output_types=["text", "audio_path"], + converter_target=adversarial_target, + ) + converter = ConverterIdentifier( + class_name="CompositeConverter", + class_module="tests.unit.memory", + supported_input_types=["text", "image_path"], + supported_output_types=["text", "audio_path"], + sub_converter=nested_converter, + ) + scorer = ScorerIdentifier( + class_name="TestScorer", + class_module="tests.unit.memory", + scorer_type="true_false", + score_aggregator="AND_", + prompt_target=adversarial_target, + ) + technique_seed = SeedIdentifier( + class_name="TechniqueSeed", + class_module="tests.unit.memory", + value="technique seed", + value_sha256="technique-sha", + data_type="text", + dataset_name="techniques", + is_general_technique=True, + ) + dataset_seed = SeedIdentifier( + class_name="DatasetSeed", + class_module="tests.unit.memory", + value="dataset seed", + value_sha256="dataset-sha", + data_type="text", + dataset_name="dataset-a", + is_general_technique=False, + ) + attack = AttackIdentifier( + class_name="TestAttack", + class_module="tests.unit.memory", + adversarial_system_prompt="system prompt", + adversarial_seed_prompt="seed prompt", + objective_target=objective_target, + adversarial_chat=adversarial_target, + objective_scorer=scorer, + request_converters=[converter], + ) + technique = AttackTechniqueIdentifier( + class_name="TestTechnique", + class_module="tests.unit.memory", + attack=attack, + technique_seeds=[technique_seed], + ) + atomic_attack = AtomicAttackIdentifier( + class_name="AtomicAttack", + class_module="tests.unit.memory", + attack_technique=technique, + seed_identifiers=[technique_seed, dataset_seed], + ) + scenario = ScenarioIdentifier( + class_name="TestScenario", + class_module="tests.unit.memory", + version=2, + techniques=["TestTechnique", "OtherTechnique"], + datasets=["dataset-a", "dataset-b"], + objective_target=objective_target, + objective_scorer=scorer, + ) + + with closing(sqlite_instance.get_session()) as session: + sqlite_instance._persist_identifier(session=session, identifier=atomic_attack) + sqlite_instance._persist_identifier(session=session, identifier=scenario) + session.commit() + + return IdentifierGraph( + objective_target=objective_target, + adversarial_target=adversarial_target, + nested_converter=nested_converter, + converter=converter, + scorer=scorer, + scenario=scenario, + technique_seed=technique_seed, + dataset_seed=dataset_seed, + attack=attack, + technique=technique, + atomic_attack=atomic_attack, + ) + + +def test_get_target_identifiers_by_hash_and_promoted_field(sqlite_instance: MemoryInterface) -> None: + target = TargetIdentifier( + class_name="TestTarget", + class_module="tests.unit.memory", + endpoint="https://example.test", + model_name="test-model", + supported_auth_modes=["api_key", "identity"], + ) + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id="identifier-query", target_identifier=target) + ) + + identifiers = sqlite_instance.get_target_identifiers( + identifier_hashes=[target.hash], + model_name="test-model", + supported_auth_modes=["identity", "api_key"], + ) + + assert identifiers == [target] + assert isinstance(identifiers[0], TargetIdentifier) + assert sqlite_instance.get_target_identifiers(supported_auth_modes=["api_key"]) == [] + assert sqlite_instance.get_target_identifiers(supported_auth_modes=["API_KEY", "identity"]) == [] + + +def test_get_identifiers_reconstructs_each_typed_graph( + sqlite_instance: MemoryInterface, identifier_graph: IdentifierGraph +) -> None: + queries = [ + ( + sqlite_instance.get_target_identifiers( + identifier_hashes=[identifier_graph.objective_target.hash], + endpoint="https://objective.test", + temperature=0.5, + ), + identifier_graph.objective_target, + ), + ( + sqlite_instance.get_converter_identifiers( + supported_input_types=["image_path", "text"], + supported_output_types=["audio_path", "text"], + sub_converter_hash=identifier_graph.nested_converter.hash, + ), + identifier_graph.converter, + ), + ( + sqlite_instance.get_scorer_identifiers( + scorer_type="true_false", + score_aggregator="AND_", + prompt_target_hash=identifier_graph.adversarial_target.hash, + ), + identifier_graph.scorer, + ), + ( + sqlite_instance.get_scenario_identifiers( + version=2, + techniques=["OtherTechnique", "TestTechnique"], + datasets=["dataset-b", "dataset-a"], + objective_target_hash=identifier_graph.objective_target.hash, + objective_scorer_hash=identifier_graph.scorer.hash, + ), + identifier_graph.scenario, + ), + ( + sqlite_instance.get_seed_identifiers( + value_sha256="dataset-sha", + data_type="text", + dataset_name="dataset-a", + is_general_technique=False, + ), + identifier_graph.dataset_seed, + ), + ( + sqlite_instance.get_attack_identifiers( + adversarial_system_prompt="system prompt", + adversarial_seed_prompt="seed prompt", + objective_target_hash=identifier_graph.objective_target.hash, + adversarial_chat_hash=identifier_graph.adversarial_target.hash, + objective_scorer_hash=identifier_graph.scorer.hash, + ), + identifier_graph.attack, + ), + ( + sqlite_instance.get_attack_technique_identifiers( + attack_identifier_hash=identifier_graph.attack.hash, + ), + identifier_graph.technique, + ), + ( + sqlite_instance.get_atomic_attack_identifiers( + attack_technique_identifier_hash=identifier_graph.technique.hash, + ), + identifier_graph.atomic_attack, + ), + ] + + for identifiers, expected in queries: + assert identifiers == [expected] + assert type(identifiers[0]) is type(expected) + + +def test_get_identifiers_common_filters_and_result_semantics( + sqlite_instance: MemoryInterface, identifier_graph: IdentifierGraph +) -> None: + targets = sqlite_instance.get_target_identifiers() + assert [identifier.hash for identifier in targets] == sorted( + [identifier_graph.objective_target.hash, identifier_graph.adversarial_target.hash] + ) + + duplicate_hashes = [identifier_graph.objective_target.hash, identifier_graph.objective_target.hash] + assert sqlite_instance.get_target_identifiers(identifier_hashes=duplicate_hashes) == [ + identifier_graph.objective_target + ] + assert sqlite_instance.get_target_identifiers(identifier_hashes=[]) == [] + assert sqlite_instance.get_target_identifiers(class_name="missing") == [] + + +def test_get_identifiers_rejects_missing_identifier_json(sqlite_instance: MemoryInterface) -> None: + identifier_hash = "0" * 64 + sqlite_instance._insert_entry(TargetIdentifierEntry(hash=identifier_hash, identifier_json=None)) + + with pytest.raises(ValueError, match="has no identifier JSON"): + sqlite_instance.get_target_identifiers(identifier_hashes=[identifier_hash]) + + +def test_get_identifiers_rejects_hash_mismatch(sqlite_instance: MemoryInterface) -> None: + target = TargetIdentifier(class_name="Target", class_module="tests.unit.memory") + identifier_hash = "f" * 64 + sqlite_instance._insert_entry(TargetIdentifierEntry(hash=identifier_hash, identifier_json=target.model_dump())) + + with pytest.raises(ValueError, match="does not match its stored JSON hash"): + sqlite_instance.get_target_identifiers(identifier_hashes=[identifier_hash]) 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..60ad1f0eb4 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 MappedColumn, 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,194 @@ def test_load_identifier_injects_pyrit_version(): assert loaded.pyrit_version == "9.9.9" +def test_scorer_identifier_entry_constructs_hash_only_sub_scorer_edges(): + 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 [edge.child_hash for edge in entry.sub_scorers] == [nested.hash, leaf.hash] + assert all(edge.child is None for edge in entry.sub_scorers) + + +@pytest.mark.parametrize( + ("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 + assert all(spec.edge_child_hash_attr for spec in entry_type.CHILD_RELATIONSHIP_SPECS.values()) + + +def test_identifier_child_relationships_delete_orphans() -> None: + for mapper in Base.registry.mappers: + entry_type = mapper.class_ + if not issubclass(entry_type, ComponentIdentifierEntry): + continue + + for field_name, spec in entry_type.CHILD_RELATIONSHIP_SPECS.items(): + child_relationship = mapper.relationships[spec.relationship_name] + assert child_relationship.uselist, f"{entry_type.__name__}.{field_name} must be a collection" + assert "delete-orphan" in child_relationship.cascade, ( + f"{entry_type.__name__}.{spec.relationship_name} must use cascade='all, delete-orphan'" + ) + + +@pytest.mark.parametrize( + ("identifier_type", "entry_type"), + [ + (TargetIdentifier, TargetIdentifierEntry), + (ConverterIdentifier, ConverterIdentifierEntry), + (ScorerIdentifier, ScorerIdentifierEntry), + (ScenarioIdentifier, ScenarioIdentifierEntry), + (SeedIdentifier, SeedIdentifierEntry), + (AttackIdentifier, AttackIdentifierEntry), + (AttackTechniqueIdentifier, AttackTechniqueIdentifierEntry), + (AtomicAttackIdentifier, AtomicAttackIdentifierEntry), + ], +) +def test_identifier_entry_maps_promoted_scalars_to_columns( + identifier_type: type[ComponentIdentifier], + entry_type: type[ComponentIdentifierEntry[Any]], +) -> None: + shared_columns = {name for name, value in vars(ComponentIdentifierEntry).items() if isinstance(value, MappedColumn)} + mapped_scalars = set(entry_type.__table__.columns.keys()) + mapped_scalars -= shared_columns + mapped_scalars -= set(entry_type.CHILD_HASH_COLUMNS.values()) + + assert mapped_scalars == set(identifier_type.promoted_scalar_field_names()) + + +def test_identifier_entry_rejects_missing_promoted_scalar_column() -> None: + class IdentifierWithPromotedScalar(ComponentIdentifier): + promoted_value: str | None = None + + table_name = "IncompleteIdentifierEntries" + try: + with pytest.raises(TypeError, match="has no mapped column for promoted scalar field.*promoted_value"): + + class IncompleteIdentifierEntry(ComponentIdentifierEntry[IdentifierWithPromotedScalar]): + __tablename__ = table_name + __table_args__ = {"extend_existing": True} + finally: + table = Base.metadata.tables.get(table_name) + if table is not None: + Base.metadata.remove(table) + + +def test_atomic_attack_identifier_graph_persists_with_result_link() -> None: + target = TargetIdentifier(class_name="Target", class_module="pyrit.prompt_target", model_name="model") + scorer = ScorerIdentifier(class_name="Scorer", class_module="pyrit.score", scorer_type="true_false") + 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..a0cc626a54 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,547 @@ def test_conversations_migration_downgrade_restores_columns(): engine.dispose() +# ============================================================================= +# Backfill tests for identifier persistence (e5f7a9c1b3d2) +# ============================================================================= + + +def test_identifier_migrations_are_nullable_and_best_effort_with_malformed_json(): + """Malformed retained identifiers do not block upgrades and leave nullable links unset.""" + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-best-effort.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "d4e6f8a0b2c4") + connection.execute( + text('INSERT INTO "Conversations" (conversation_id, target_identifier) VALUES (:id, :value)'), + {"id": "malformed-conversation", "value": "not-json"}, + ) + score_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "ScoreEntries" ' + "(id, score_value, score_type, score_metadata, scorer_class_identifier, timestamp) " + "VALUES (:id, 'True', 'true_false', '{}', 'not-json', '2026-07-13')" + ), + {"id": score_id}, + ) + result_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "ScenarioResultEntries" ' + "(id, scenario_name, scenario_version, pyrit_version, scenario_identifier, " + "objective_target_identifier, scenario_run_state, attack_results_json, number_tries, " + "completion_time, timestamp) VALUES (:id, 'Scenario', 1, '0.10.0', 'not-json', '{}', " + "'COMPLETED', '{}', 1, '2026-07-13', '2026-07-13')" + ), + {"id": result_id}, + ) + prompt_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "PromptMemoryEntries" ' + "(id, role, conversation_id, sequence, timestamp, labels, prompt_metadata, " + "converter_identifiers, original_value_data_type, original_value, converted_value_data_type, " + "original_prompt_id) VALUES (:id, 'user', 'conversation', 0, '2026-07-13', '{}', '{}', " + "'not-json', 'text', 'prompt', 'text', :id)" + ), + {"id": prompt_id}, + ) + attack_result_id = str(uuid.uuid4()) + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, atomic_attack_identifier, executed_turns, execution_time_ms, " + "outcome, timestamp) VALUES (:id, 'conversation', 'objective', 'not-json', 0, 0, " + "'undetermined', '2026-07-13')" + ), + {"id": attack_result_id}, + ) + + command.upgrade(config, "e5f7a9c1b3d2") + + assert ( + connection.execute( + text( + 'SELECT target_identifier_hash FROM "Conversations" ' + "WHERE conversation_id = 'malformed-conversation'" + ) + ).scalar_one() + is None + ) + assert ( + connection.execute( + text('SELECT scorer_identifier_hash FROM "ScoreEntries" WHERE id = :id'), + {"id": score_id}, + ).scalar_one() + is None + ) + assert ( + connection.execute( + text('SELECT scenario_identifier_hash FROM "ScenarioResultEntries" WHERE id = :id'), + {"id": result_id}, + ).scalar_one() + is None + ) + assert connection.execute(text('SELECT COUNT(*) FROM "PromptConverterIdentifiers"')).scalar_one() == 0 + assert ( + connection.execute( + text('SELECT atomic_attack_identifier_hash FROM "AttackResultEntries" WHERE id = :id'), + {"id": attack_result_id}, + ).scalar_one() + is None + ) + + for table_name in ( + "TargetIdentifiers", + "ScorerIdentifiers", + "ScenarioIdentifiers", + "ConverterIdentifiers", + "SeedIdentifiers", + "AttackIdentifiers", + "AttackTechniqueIdentifiers", + "AtomicAttackIdentifiers", + ): + columns = {column["name"]: column for column in inspect(connection).get_columns(table_name)} + assert columns["hash"]["nullable"] is False + assert columns["class_name"]["nullable"] is True + assert columns["class_module"]["nullable"] is True + assert columns["identifier_json"]["nullable"] is True + finally: + engine.dispose() + + +def test_identifier_migration_reuses_edges_and_rolls_back_conflicting_row(): + """Repeated graphs reuse edges while a conflicting row is isolated in its savepoint.""" + child_hash = "a" * 64 + conflicting_child_hash = "b" * 64 + parent_hash = "c" * 64 + healthy_hash = "d" * 64 + child = {"hash": child_hash, "class_name": "Child", "class_module": "test"} + parent = { + "hash": parent_hash, + "class_name": "Parent", + "class_module": "test", + "children": {"targets": [child]}, + } + conflicting_parent = { + **parent, + "children": {"targets": [{"hash": conflicting_child_hash, "class_name": "OtherChild", "class_module": "test"}]}, + } + healthy = {"hash": healthy_hash, "class_name": "Healthy", "class_module": "test"} + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-row-savepoints.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "d4e6f8a0b2c4") + insert_statement = text( + 'INSERT INTO "Conversations" (conversation_id, target_identifier) VALUES (:id, :identifier)' + ) + for conversation_id, identifier in ( + ("01-original", parent), + ("02-duplicate", parent), + ("03-conflict", conflicting_parent), + ("04-healthy", healthy), + ): + connection.execute( + insert_statement, + {"id": conversation_id, "identifier": json.dumps(identifier)}, + ) + + command.upgrade(config, "e5f7a9c1b3d2") + + links = connection.execute( + text('SELECT conversation_id, target_identifier_hash FROM "Conversations" ORDER BY conversation_id') + ).fetchall() + target_hashes = set(connection.execute(text('SELECT hash FROM "TargetIdentifiers"')).scalars()) + edges = connection.execute( + text('SELECT parent_hash, position, child_hash FROM "TargetIdentifierChildren"') + ).fetchall() + + assert links == [ + ("01-original", parent_hash), + ("02-duplicate", parent_hash), + ("03-conflict", None), + ("04-healthy", healthy_hash), + ] + assert target_hashes == {child_hash, parent_hash, healthy_hash} + assert edges == [(parent_hash, 0, child_hash)] + finally: + engine.dispose() + + +def test_identifier_migration_downgrade_drops_link_constraints_and_columns(): + """Downgrade removes identifier foreign keys before removing their columns.""" + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'identifier-downgrade.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "e5f7a9c1b3d2") + command.downgrade(config, "d4e6f8a0b2c4") + + assert "target_identifier_hash" not in { + column["name"] for column in inspect(connection).get_columns("Conversations") + } + assert "TargetIdentifiers" not in set(inspect(connection).get_table_names()) + finally: + engine.dispose() + + +def test_identifier_migrations_do_not_import_domain_models(): + """Frozen identifier migrations operate on retained JSON rather than current domain models.""" + versions_dir = Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic" / "versions" + revision_names = ("e5f7a9c1b3d2_add_identifiers_tables.py",) + + for revision_name in revision_names: + source = (versions_dir / revision_name).read_text(encoding="utf-8") + assert "from pyrit.models" not in source + assert "import pyrit.models" not in source + + +def test_scorer_identifier_migration_backfills_graph_and_score_link(): + """Existing scorer JSON is materialized and linked without changing its content identity.""" + from pyrit.models import ComponentIdentifier + + 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, "d4e6f8a0b2c4") + 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, "e5f7a9c1b3d2") + + 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 +# ============================================================================= + + +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, "d4e6f8a0b2c4") + 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, "e5f7a9c1b3d2") + + 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 +# ============================================================================= + + +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, "d4e6f8a0b2c4") + 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, "e5f7a9c1b3d2") + + 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 +# ============================================================================= + + +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, "d4e6f8a0b2c4") + 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, "e5f7a9c1b3d2") + + 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"}