Skip to content

[DRAFT] Persist identifiers in DB#2164

Draft
behnam-o wants to merge 8 commits into
microsoft:mainfrom
behnam-o:persist-identifiers
Draft

[DRAFT] Persist identifiers in DB#2164
behnam-o wants to merge 8 commits into
microsoft:mainfrom
behnam-o:persist-identifiers

Conversation

@behnam-o

Copy link
Copy Markdown
Contributor

Problem

This is phase 7 of this proposal: PyRIT Unified Registry — Design & Rationale

this draft uses TargetIdentifier and their relationship with Conversation as example, if we concensus on the pattern, this will be extended to all identifiers.

Target identifiers weren't persisted as first-class, queryable data. Their configuration (endpoint, model, sampling params) and their composition (multi-targets wrapping inner targets) were only available inline, with no deduplication and no way to query across them.

Solution

Persist TargetIdentifiers as content-addressed rows keyed by their content hash, so an identical identifier reused across many conversations maps to a single immutable row.

Mechanisms:

  • Content-addressed value store — the full identifier dump (identifier_json) is the immutable source of truth; frequently-queried fields are surfaced as denormalized columns (indexes into the blob, not a second source of truth).
  • from_domain_model seam — a uniform, per-entry converter defining how a domain model becomes a row, enforced on concrete subclasses at class-definition time.
  • Recursive persistence with a child graph — multi-target composition is stored as an ordered parent→child edge table, materialized by a recursive write path that mirrors the existing "persist dependencies first" contract, with idempotent insert-if-absent semantics.

Why draft

Raising as a draft to validate the pattern (DomainBackedEntry + content-addressed store + recursive child graph) on target identifiers first. If it looks good, the same approach extends cleanly to the other identifier types.

@rlundeen2

Copy link
Copy Markdown
Contributor

I wonder if we can figure out how to not repeat the promoted values 3x per identifier. Here is a way that copilot brainstormed that shows promise:

The promotion metadata already lives on the identifier (ComponentIdentifier._promoted_param_fields()), so TargetIdentifierEntry could read the field set off the identifier instead of restating it. There are two coupling points to collapse — the value mapping (from_domain_model) and the column set — and the PR's own __init_subclass__ pattern is a good fit for enforcing the second. (The Alembic backfill also restates the list, but migrations are frozen snapshots and should stay hand-written.)

1. Expose a public accessor on the domain model

_promoted_param_fields() exists but is private. A thin public seam on ComponentIdentifier lets memory read the projection without crossing a privacy boundary:

# component_identifier.py
@classmethod
def promoted_scalar_field_names(cls) -> tuple[str, ...]:
    """Names of this identifier's promoted scalar (param) fields — the DB-column projection."""
    return cls._promoted_param_fields()

def promoted_scalar_values(self) -> dict[str, Any]:
    """This identifier's promoted scalar fields as ``{name: value}`` (children/targets excluded)."""
    return {name: getattr(self, name) for name in self._promoted_param_fields()}

For TargetIdentifier this yields exactly endpoint, model_name, underlying_model_name, temperature, top_p, max_requests_per_minute, supported_auth_modes — no hand-maintained list.

2. Make from_domain_model generic in the base

Drop the hand-written mapping in TargetIdentifierEntry; the base handles every subclass:

# ComponentIdentifierEntry
@classmethod
def from_domain_model(cls, domain_model: T) -> Self:
    entry = cls(
        hash=domain_model.hash,
        class_name=domain_model.class_name,
        class_module=domain_model.class_module,
        identifier_json=domain_model.model_dump(),
        pyrit_version=domain_model.pyrit_version,
    )
    for name, value in domain_model.promoted_scalar_values().items():
        setattr(entry, name, value)  # each promoted scalar → its mapped column
    return entry

TargetIdentifierEntry then declares only its columns — no from_domain_model override.

3. Turn column drift into an import-time error (reuses the pattern already in this PR)

This PR already validates entries at class-definition time via DomainBackedEntry.__init_subclass__. The same hook can assert columns == promoted scalar fields, recovering the bound identifier type from the generic base:

import typing

def _bound_identifier_type(cls) -> type[ComponentIdentifier] | None:
    for base in getattr(cls, "__orig_bases__", ()):
        if typing.get_origin(base) is ComponentIdentifierEntry:
            (arg,) = typing.get_args(base)
            return arg if isinstance(arg, type) else None
    return None

# in __init_subclass__, for concrete subclasses:
identifier_type = _bound_identifier_type(cls)
if identifier_type is not None:
    shared = set(ComponentIdentifierEntry.__mapper__.columns.keys())  # hash/class_name/...
    own_columns = set(cls.__table__.columns.keys()) - shared
    promoted = set(identifier_type.promoted_scalar_field_names())
    if own_columns != promoted:
        raise TypeError(
            f"{cls.__name__} columns {sorted(own_columns)} do not match "
            f"{identifier_type.__name__} promoted scalar fields {sorted(promoted)}; "
            "add/remove the mapped column to match the identifier."
        )

Now adding a promoted scalar to TargetIdentifier and forgetting the column fails at import, with a message that says exactly what to add. This also closes the one hole in step 2: setattr on a missing column would otherwise silently set a non-persisted Python attr.

Deliberately not auto-generating the columns

Generating the mapped_columns from the pydantic annotations is possible but probably not worth it: you'd lose the static TargetIdentifierEntry.model_name attributes that IDEs, the type-checker, and Alembic autogenerate rely on. Keeping columns explicit + a class-definition-time parity check feels like the sweet spot: explicit schema, single source of truth for which fields, mechanical enforcement.

Net effect: the promoted-field list lives in exactly one authoritative place (the identifier), the ORM derives values from it, and the class-definition-time check keeps the column DDL honest — collapsing the 3× repetition down to the identifier plus the (necessarily frozen) migration.

@@ -459,6 +462,13 @@ def _insert_conversation(self, *, conversation: Conversation) -> None:
try:
existing = session.get(ConversationEntry, conversation.conversation_id)
if existing is None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know this is draft, but I'd also add queries to the same PR. I think it makes testability easier even at the expense of bigger PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants