From 269b4f93f7f55bd8be49bff59f27131b0c5ec6aa Mon Sep 17 00:00:00 2001 From: behnamousat Date: Fri, 10 Jul 2026 15:41:26 -0700 Subject: [PATCH] Persist target identifiers as a deduped, content-addressed store --- ...f7a9c1b3d2_add_target_identifiers_table.py | 203 ++++++++++++++++++ pyrit/memory/memory_interface.py | 60 +++++- pyrit/memory/memory_models.py | 194 ++++++++++++++++- .../test_interface_prompts.py | 96 +++++++++ 4 files changed, 551 insertions(+), 2 deletions(-) create mode 100644 pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py diff --git a/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py new file mode 100644 index 0000000000..f64af89ef6 --- /dev/null +++ b/pyrit/memory/alembic/versions/e5f7a9c1b3d2_add_target_identifiers_table.py @@ -0,0 +1,203 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Introduce the TargetIdentifiers table and reference it from Conversations. + +Phase 1 (dual-write) of storing component identifiers as first-class, +content-addressed rows. Creates ``TargetIdentifiers`` (one row per distinct +target identifier, keyed by its content ``hash``, with promoted scalar query +columns) and ``TargetIdentifierChildren`` (a self-referential pivot mapping a +multi-target to its inner target identifiers), adds a nullable +``target_identifier_hash`` foreign key to ``Conversations``, and backfills all +three from the existing ``Conversations.target_identifier`` JSON column. The JSON +column is retained (reads still come from it), so this migration is purely +additive. + +Revision ID: e5f7a9c1b3d2 +Revises: d4e6f8a0b2c4 +Create Date: 2026-07-10 12:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "e5f7a9c1b3d2" +down_revision: str | None = "d4e6f8a0b2c4" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_table( + "TargetIdentifiers", + sa.Column("hash", sa.String(64), primary_key=True, nullable=False), + sa.Column("class_name", sa.String(), nullable=False), + sa.Column("class_module", sa.String(), nullable=False), + sa.Column("identifier_json", sa.JSON(), nullable=False), + sa.Column("endpoint", sa.String(), nullable=True), + sa.Column("model_name", sa.String(), nullable=True), + sa.Column("underlying_model_name", sa.String(), nullable=True), + sa.Column("temperature", sa.Float(), nullable=True), + sa.Column("top_p", sa.Float(), nullable=True), + sa.Column("max_requests_per_minute", sa.Integer(), nullable=True), + sa.Column("supported_auth_modes", sa.JSON(), nullable=True), + sa.Column("pyrit_version", sa.String(), nullable=True), + ) + + # Self-referential pivot mapping a multi-target to its inner target identifiers. + # Both endpoints are content hashes into TargetIdentifiers; ``position`` preserves + # the parent's ``targets`` list order. Named FK constraints for SQL Server / batch + # portability. + op.create_table( + "TargetIdentifierChildren", + sa.Column("parent_hash", sa.String(64), nullable=False), + sa.Column("position", sa.Integer(), nullable=False), + sa.Column("child_hash", sa.String(64), nullable=False), + sa.PrimaryKeyConstraint("parent_hash", "position"), + sa.ForeignKeyConstraint( + ["parent_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_parent_hash" + ), + sa.ForeignKeyConstraint( + ["child_hash"], ["TargetIdentifiers.hash"], name="fk_target_identifier_children_child_hash" + ), + ) + + # Batch op for SQLite portability (no ALTER TABLE ADD FOREIGN KEY on SQLite). + # The FK constraint must be named explicitly: Alembic batch mode rejects an + # unnamed constraint. + with op.batch_alter_table("Conversations") as batch_op: + batch_op.add_column(sa.Column("target_identifier_hash", sa.String(64), nullable=True)) + batch_op.create_foreign_key( + "fk_conversations_target_identifier_hash", + "TargetIdentifiers", + ["target_identifier_hash"], + ["hash"], + ) + + _backfill_target_identifiers() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + with op.batch_alter_table("Conversations") as batch_op: + batch_op.drop_column("target_identifier_hash") + # Drop the child edge table before its referenced parent table. + op.drop_table("TargetIdentifierChildren") + op.drop_table("TargetIdentifiers") + + +def _backfill_target_identifiers() -> None: + """ + Populate ``TargetIdentifiers`` / ``TargetIdentifierChildren`` and set + ``Conversations.target_identifier_hash``. + + For every ``Conversations`` row with a non-null ``target_identifier`` JSON, + reconstruct the ``TargetIdentifier`` (recomputing its content hash), insert the + deduped ``TargetIdentifiers`` row if absent -- recursing into any inner + ``targets`` first so the child edge foreign keys resolve -- record the + ``parent_hash -> child_hash`` edges, and point the conversation's + ``target_identifier_hash`` at the top-level row. Idempotent: hashes already present + are not re-inserted. Rows whose stored target cannot be reconstructed are logged and + skipped rather than aborting the upgrade. + """ + from pyrit.models import TargetIdentifier + + bind = op.get_bind() + rows = bind.execute( + sa.text('SELECT conversation_id, target_identifier FROM "Conversations" WHERE target_identifier IS NOT NULL') + ).fetchall() + + existing_hashes = {row[0] for row in bind.execute(sa.text('SELECT hash FROM "TargetIdentifiers"')).fetchall()} + + insert_stmt = sa.text( + 'INSERT INTO "TargetIdentifiers" ' + "(hash, class_name, class_module, identifier_json, endpoint, model_name, " + "underlying_model_name, temperature, top_p, max_requests_per_minute, " + "supported_auth_modes, pyrit_version) " + "VALUES (:hash, :class_name, :class_module, :identifier_json, :endpoint, :model_name, " + ":underlying_model_name, :temperature, :top_p, :max_requests_per_minute, " + ":supported_auth_modes, :pyrit_version)" + ) + edge_stmt = sa.text( + 'INSERT INTO "TargetIdentifierChildren" (parent_hash, position, child_hash) ' + "VALUES (:parent_hash, :position, :child_hash)" + ) + update_stmt = sa.text('UPDATE "Conversations" SET target_identifier_hash = :hash WHERE conversation_id = :cid') + + def _insert_target(identifier: TargetIdentifier) -> int: + """ + Insert ``identifier`` and its descendants if absent; record child edges. + + Returns: + int: The number of new ``TargetIdentifiers`` rows inserted (self + descendants). + """ + target_hash = identifier.hash + if target_hash in existing_hashes: + return 0 + # Children first so the parent's edge foreign keys resolve. + inserted = 0 + for child in identifier.targets: + inserted += _insert_target(child) + bind.execute( + insert_stmt, + { + "hash": target_hash, + "class_name": identifier.class_name, + "class_module": identifier.class_module, + "identifier_json": json.dumps(identifier.model_dump(), sort_keys=True), + "endpoint": identifier.endpoint, + "model_name": identifier.model_name, + "underlying_model_name": identifier.underlying_model_name, + "temperature": identifier.temperature, + "top_p": identifier.top_p, + "max_requests_per_minute": identifier.max_requests_per_minute, + "supported_auth_modes": ( + json.dumps(identifier.supported_auth_modes) if identifier.supported_auth_modes is not None else None + ), + "pyrit_version": identifier.pyrit_version, + }, + ) + existing_hashes.add(target_hash) + inserted += 1 + for position, child in enumerate(identifier.targets): + bind.execute(edge_stmt, {"parent_hash": target_hash, "position": position, "child_hash": child.hash}) + return inserted + + inserted = 0 + linked = 0 + skipped = 0 + for conversation_id, raw_target in rows: + stored = json.loads(raw_target) if isinstance(raw_target, str) else raw_target + if not stored: + continue + try: + identifier = TargetIdentifier.model_validate(stored) + except Exception: + skipped += 1 + logger.warning( + f"TargetIdentifiers backfill: could not reconstruct target for " + f"conversation_id {conversation_id!r}; skipping." + ) + continue + + inserted += _insert_target(identifier) + bind.execute(update_stmt, {"hash": identifier.hash, "cid": conversation_id}) + linked += 1 + + if inserted or linked or skipped: + logger.info( + f"TargetIdentifiers backfill: inserted {inserted} identifier row(s); " + f"linked {linked} conversation(s); skipped {skipped}." + ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 9e2bf380af..7ff7943c21 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -14,7 +14,7 @@ from sqlalchemy import MetaData, and_, not_, or_, select from sqlalchemy.engine.base import Engine -from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.orm.attributes import InstrumentedAttribute if TYPE_CHECKING: @@ -29,6 +29,8 @@ ScenarioResultEntry, ScoreEntry, SeedEntry, + TargetIdentifierChildEntry, + TargetIdentifierEntry, ) from pyrit.memory.storage import ( DataTypeSerializer, @@ -50,6 +52,7 @@ SeedDataset, SeedGroup, SeedType, + TargetIdentifier, group_conversation_message_pieces_by_sequence, sort_message_pieces, ) @@ -459,6 +462,13 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: try: existing = session.get(ConversationEntry, conversation.conversation_id) if existing is None: + if conversation.target_identifier is not None: + self._persist_target_identifier( + session=session, + target_identifier=TargetIdentifier.from_component_identifier( + conversation.target_identifier + ), + ) session.add(entry) elif ( entry.target_identifier is not None @@ -476,6 +486,54 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}") raise + @staticmethod + def _persist_target_identifier(*, session: Any, target_identifier: TargetIdentifier) -> None: + """ + Persist ``target_identifier`` and its inner targets as content-addressed rows. + + Mirrors the conversation-registration contract ("ensure dependencies exist + first"): every inner target in ``target_identifier.targets`` is persisted + recursively before this identifier's own row, so the ``TargetIdentifierChildren`` + edges (which foreign-key both endpoints into ``TargetIdentifiers``) resolve. + Identifier rows are immutable and keyed by their content hash, so this is an + idempotent insert-if-absent: an identical target -- or a shared inner target -- + reused across many conversations maps to a single row. Runs inside the caller's + session so every row is flushed before the FK-bearing ``ConversationEntry`` is + added. + + If the row already exists it was fully persisted before (children and edges + included, since rows are immutable), so this returns early. Otherwise the row and + its child edges are inserted inside a savepoint: a concurrent writer that inserts + the same hash first surfaces as an ``IntegrityError`` that rolls back only this + insert (not the surrounding write) and is then treated as a no-op -- the row + exists either way. + + Args: + session (Any): The active SQLAlchemy session (the caller's transaction). + target_identifier (TargetIdentifier): The target identifier to persist. + """ + if session.get(TargetIdentifierEntry, target_identifier.hash) is not None: + return + # Children first, so the parent's edge foreign keys resolve on flush. + for inner_target in target_identifier.targets: + MemoryInterface._persist_target_identifier(session=session, target_identifier=inner_target) + try: + with session.begin_nested(): + session.add(TargetIdentifierEntry.from_domain_model(target_identifier)) + for position, inner_target in enumerate(target_identifier.targets): + session.add( + TargetIdentifierChildEntry( + parent_hash=target_identifier.hash, + position=position, + child_hash=inner_target.hash, + ) + ) + session.flush() + except IntegrityError: + # A racing writer inserted the same content-addressed row; immutable + # rows make this a safe no-op. + pass + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a7cfdbf552..04a0fb568e 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -4,9 +4,10 @@ import json import logging import uuid +from abc import abstractmethod from collections.abc import Sequence from datetime import datetime, timezone -from typing import Any, Literal +from typing import Any, Generic, Literal, TypeVar from pydantic import BaseModel, ConfigDict from sqlalchemy import ( @@ -29,6 +30,7 @@ relationship, ) from sqlalchemy.types import Uuid +from typing_extensions import Self import pyrit from pyrit.common.utils import to_sha256 @@ -56,6 +58,7 @@ SeedPrompt, SeedSimulatedConversation, SeedType, + TargetIdentifier, ) logger = logging.getLogger(__name__) @@ -353,6 +356,185 @@ def __str__(self) -> str: return f"{self.role}: {self.converted_value}" +TDomain = TypeVar("TDomain") + + +class DomainBackedEntry(Base, Generic[TDomain]): + """ + Mixin marking a DB entry as the persistence representation of a domain model. + + Every ``*Entry`` in this module mirrors a domain model (``PromptMemoryEntry`` and + ``MessagePiece``, ``ScoreEntry`` and ``Score``, ``TargetIdentifierEntry`` and + ``TargetIdentifier``, and so on). ``from_domain_model`` is the single, uniform seam + that converts a domain model into an unsaved row, so the domain-to-DB direction has + one well-known name across every entry that adopts this base. + """ + + __abstract__ = True + + @classmethod + @abstractmethod + def from_domain_model(cls, domain_model: TDomain) -> Self: + """ + Build an unsaved entry row from its domain model. + + Args: + domain_model (TDomain): The domain model this entry persists. + + Returns: + Self: A new, unsaved row. + """ + + def __init_subclass__(cls, **kwargs: Any) -> None: + """ + Reject concrete subclasses that do not implement ``from_domain_model``. + + The SQLAlchemy declarative ``Base`` is not an ``ABCMeta``, so a bare + ``@abstractmethod`` would not stop a concrete entry from omitting the converter. + This fires at class-definition time so a dev who forgets is told immediately. + SQLAlchemy abstract/intermediate mapped classes (``__abstract__ = True``) are + skipped so they can leave the method abstract for their concrete subclasses. + + Raises: + TypeError: If a concrete (non-``__abstract__``) subclass leaves + ``from_domain_model`` abstract. + """ + super().__init_subclass__(**kwargs) + if cls.__dict__.get("__abstract__", False): + return + method = getattr(cls, "from_domain_model", None) + if method is None or getattr(method, "__isabstractmethod__", False): + raise TypeError( + f"{cls.__name__} inherits DomainBackedEntry but does not implement " + "from_domain_model(...); every concrete entry must define how its " + "domain model is converted into a row." + ) + + +T = TypeVar("T", bound=ComponentIdentifier) + + +class ComponentIdentifierEntry(DomainBackedEntry[T]): + """ + Abstract base for tables that persist a ``ComponentIdentifier`` projection. + + Mirrors the identifier class hierarchy: concrete identifier tables inherit the + shared, always-populated columns the way ``TargetIdentifier`` inherits + ``ComponentIdentifier``. The content ``hash`` is the natural, dedupable primary + key, and ``identifier_json`` holds the full flat dump for lossless + reconstruction (children stay inline). Rows are immutable — the same content + always maps to the same hash, so a given identifier reused across rows is + stored once. + + Subclasses declare their promoted query columns and implement the + ``DomainBackedEntry.from_domain_model`` seam to map their strongly-typed identifier + projection onto the shared columns plus those promoted columns. The shared columns + are built once, here, so subclasses never repeat that logic. + """ + + __abstract__ = True + + #: Content-addressed identity — the same value as ``ComponentIdentifier.hash``. + #: SHA256 hex digest is 64 chars; bounded for SQL Server key/index compatibility. + hash: Mapped[str] = mapped_column(String(64), primary_key=True) + class_name: Mapped[str] = mapped_column(String, nullable=False) + class_module: Mapped[str] = mapped_column(String, nullable=False) + #: Full flat ``model_dump()`` of the identifier. Source of truth on reload. + identifier_json: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) + #: Version that first wrote this content-addressed row. Nullable for backwards + #: compatibility with existing databases. + pyrit_version: Mapped[str | None] = mapped_column(String, nullable=True) + + +class TargetIdentifierEntry(ComponentIdentifierEntry[TargetIdentifier]): + """ + Content-addressed store of ``TargetIdentifier`` projections, deduped by hash. + + Populated as a side effect of registering a conversation (see + ``MemoryInterface._persist_target_identifier``). ``ConversationEntry`` references a + row here via ``target_identifier_hash``. The promoted scalar columns (``endpoint`` / + ``model_name`` / ``underlying_model_name`` / ``temperature`` / ``top_p`` / + ``max_requests_per_minute`` / ``supported_auth_modes``) are surfaced for querying; + ``identifier_json`` remains the source of truth on reload. Inner targets of a + multi-target are linked via ``TargetIdentifierChildren`` (and also live inline in + ``identifier_json``). + """ + + __tablename__ = "TargetIdentifiers" + __table_args__ = {"extend_existing": True} + + endpoint: Mapped[str | None] = mapped_column(String, nullable=True) + model_name: Mapped[str | None] = mapped_column(String, nullable=True) + underlying_model_name: Mapped[str | None] = mapped_column(String, nullable=True) + temperature: Mapped[float | None] = mapped_column(Float, nullable=True) + top_p: Mapped[float | None] = mapped_column(Float, nullable=True) + max_requests_per_minute: Mapped[int | None] = mapped_column(INTEGER, nullable=True) + supported_auth_modes: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) + + @classmethod + def from_domain_model(cls, domain_model: TargetIdentifier) -> Self: + """ + Build a row from a ``TargetIdentifier``. + + Maps the shared columns plus the promoted scalar query columns; the full flat + dump (including any inner ``targets``) is stored in ``identifier_json``. Inner + target edges are written separately by + ``MemoryInterface._persist_target_identifier``. + + Args: + domain_model (TargetIdentifier): The target identifier to persist. + + Returns: + Self: A new, unsaved row keyed by the identifier's content hash. + """ + return cls( + hash=domain_model.hash, + class_name=domain_model.class_name, + class_module=domain_model.class_module, + identifier_json=domain_model.model_dump(), + pyrit_version=domain_model.pyrit_version, + endpoint=domain_model.endpoint, + model_name=domain_model.model_name, + underlying_model_name=domain_model.underlying_model_name, + temperature=domain_model.temperature, + top_p=domain_model.top_p, + max_requests_per_minute=domain_model.max_requests_per_minute, + supported_auth_modes=domain_model.supported_auth_modes, + ) + + +class TargetIdentifierChildEntry(Base): + """ + Ordered edge rows linking a multi-target ``TargetIdentifierEntry`` to its inner + target identifiers. + + A multi-target (e.g. ``RoundRobinTarget``) owns a ``targets`` list; each inner + target is itself a content-addressed ``TargetIdentifiers`` row, and one edge row + here maps ``parent_hash -> child_hash`` at a given ``position`` (the child's index + in the parent's ``targets`` list). Both endpoints are hashes into + ``TargetIdentifiers``, so an inner target shared across parents dedupes to a single + row and is merely referenced here. This is a query index over target composition; + ``TargetIdentifierEntry.identifier_json`` remains the source of truth for + reconstruction (inner targets are stored inline there too). + + Materialized by ``MemoryInterface._persist_target_identifier`` while persisting a + target and its children; it is never built from a single domain model, so it is a + plain ``Base`` row rather than a ``DomainBackedEntry``. + """ + + __tablename__ = "TargetIdentifierChildren" + __table_args__ = {"extend_existing": True} + + parent_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), primary_key=True + ) + #: Zero-based index of the child within the parent's ``targets`` list. + position: Mapped[int] = mapped_column(INTEGER, primary_key=True) + child_hash: Mapped[str] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=False + ) + + class ConversationEntry(Base): """ Conversation-scoped metadata, persisted once per ``conversation_id``. @@ -362,6 +544,10 @@ class ConversationEntry(Base): row. The target is captured once when the conversation's pieces are written and read back via ``MemoryInterface._get_conversation`` (it is not stamped onto individual pieces). + + The target is dual-written: the full identifier stays in the ``target_identifier`` + JSON column (still the read source), and ``target_identifier_hash`` references the + deduped ``TargetIdentifierEntry`` row keyed by the identifier's content hash. """ __tablename__ = "Conversations" @@ -369,6 +555,11 @@ class ConversationEntry(Base): conversation_id = mapped_column(String(36), primary_key=True, nullable=False) target_identifier: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True) + #: Foreign key to the content-addressed ``TargetIdentifiers`` row. Nullable: + #: a conversation may be registered without a known target. + target_identifier_hash: Mapped[str | None] = mapped_column( + String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True + ) # Version of PyRIT used when this entry was created. Nullable for backwards # compatibility with existing databases. @@ -383,6 +574,7 @@ def __init__(self, *, conversation: Conversation) -> None: """ self.conversation_id = conversation.conversation_id self.target_identifier = conversation.target_identifier.model_dump() if conversation.target_identifier else None + self.target_identifier_hash = conversation.target_identifier.hash if conversation.target_identifier else None self.pyrit_version = pyrit.__version__ def get_conversation(self) -> Conversation: diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index b3860519e3..3a50b09ceb 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -679,6 +679,102 @@ def test_add_conversation_to_memory_different_target_reregister_raises(sqlite_in assert metadata.target_identifier.hash == target_a.hash +def test_target_identifier_dual_write_reconstruction_is_equivalent(sqlite_instance: MemoryInterface): + # Phase 1 dual-write invariant: reconstructing the target from the ConversationEntry + # JSON column must be identical to reconstructing it from the deduped + # TargetIdentifierEntry.identifier_json row, and the stored PK must match the + # recomputed content hash. This is the safety net that lets phase 2 flip reads to + # the FK path. + from pyrit.memory.memory_models import ConversationEntry, TargetIdentifierEntry + + target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://api.openai.com", "model_name": "gpt-4"}, + ) + conversation_id = "conv-dualwrite" + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id=conversation_id, target_identifier=target) + ) + + conv_entry = sqlite_instance._query_entries( + ConversationEntry, conditions=ConversationEntry.conversation_id == conversation_id + )[0] + id_entry = sqlite_instance._query_entries( + TargetIdentifierEntry, conditions=TargetIdentifierEntry.hash == target.hash + )[0] + + from_conversation_json = ComponentIdentifier.model_validate(conv_entry.target_identifier) + from_identifier_row = ComponentIdentifier.model_validate(id_entry.identifier_json) + + # Both reconstructions agree (equality is content-hash based) and match the FK / PK. + assert from_conversation_json == from_identifier_row + assert from_conversation_json.hash == target.hash + assert id_entry.hash == target.hash + assert conv_entry.target_identifier_hash == target.hash + # Promoted columns are surfaced from params for querying. + assert id_entry.endpoint == "https://api.openai.com" + assert id_entry.model_name == "gpt-4" + + +def test_target_identifier_row_is_deduped_across_conversations(sqlite_instance: MemoryInterface): + # The same target reused across conversations is content-addressed, so it is stored + # once: two conversations with an identical target share a single TargetIdentifiers row. + from pyrit.memory.memory_models import TargetIdentifierEntry + + target = ComponentIdentifier( + class_name="OpenAIChatTarget", class_module="pyrit.prompt_target", params={"endpoint": "shared"} + ) + for cid in ("conv-dedup-a", "conv-dedup-b"): + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id=cid, target_identifier=target) + ) + + rows = sqlite_instance._query_entries(TargetIdentifierEntry, conditions=TargetIdentifierEntry.hash == target.hash) + assert len(rows) == 1 + + +def test_target_identifier_persists_inner_targets_and_edges(sqlite_instance: MemoryInterface): + # A multi-target's inner targets are persisted as their own content-addressed rows + # and linked to the parent via ordered TargetIdentifierChildren edges. Promoted + # scalar columns are surfaced on each row for querying. + from pyrit.memory.memory_models import TargetIdentifierChildEntry, TargetIdentifierEntry + + inner_a = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://a", "model_name": "gpt-a", "temperature": 0.5}, + ) + inner_b = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"endpoint": "https://b", "model_name": "gpt-b"}, + ) + multi = ComponentIdentifier( + class_name="RoundRobinTarget", + class_module="pyrit.prompt_target", + children={"targets": [inner_a, inner_b]}, + ) + sqlite_instance.add_conversation_to_memory( + conversation=Conversation(conversation_id="conv-multi", target_identifier=multi) + ) + + id_rows = sqlite_instance._query_entries(TargetIdentifierEntry) + by_hash = {row.hash: row for row in id_rows} + # Parent + both inner targets are each stored once (content-addressed). + assert multi.hash in by_hash + assert inner_a.hash in by_hash + assert inner_b.hash in by_hash + # Promoted scalar column surfaced from the inner target's params. + assert by_hash[inner_a.hash].temperature == 0.5 + + edges = sqlite_instance._query_entries( + TargetIdentifierChildEntry, conditions=TargetIdentifierChildEntry.parent_hash == multi.hash + ) + ordered = sorted(edges, key=lambda edge: edge.position) + assert [(edge.position, edge.child_hash) for edge in ordered] == [(0, inner_a.hash), (1, inner_b.hash)] + + def test_insert_conversation_rolls_back_and_reraises_on_db_error(sqlite_instance: MemoryInterface): # A DB failure during registration rolls back the session and propagates the error # rather than leaving a half-written Conversations row.