Skip to content

[DEMONSTRATION ONLY - WILL NOT BE MERGED!] Component Config / Registry Entry: A new alternative design to Identifiers#1369

Open
bashirpartovi wants to merge 2 commits intoAzure:mainfrom
bashirpartovi:dev/bashirpartovi/id
Open

[DEMONSTRATION ONLY - WILL NOT BE MERGED!] Component Config / Registry Entry: A new alternative design to Identifiers#1369
bashirpartovi wants to merge 2 commits intoAzure:mainfrom
bashirpartovi:dev/bashirpartovi/id

Conversation

@bashirpartovi
Copy link
Contributor

@bashirpartovi bashirpartovi commented Feb 12, 2026

Simplify Component Identity: ComponentConfig + RegistryEntry

NOTE:
For demonstration purposes, this is additive, I did not remove any of the existing Identifier functionality. Components implement both Identifiable and Configurable for side-by-side comparison. For RegistryEntry, the metadata classes (ScenarioMetadata, InitializerMetadata) are swapped directly since they are ephemeral and never persisted to storage.

Context

I was looking at the Identifier class and realized it's overly complex for what it does. It handles instance identity, registry metadata, hash computation with field exclusion annotations, serialization with storage truncation, and more, all in one class. It works, but the coupling it creates makes the codebase harder to change than it should be.

I started asking what if we split this into two small, focused types instead? That's what this PR does. Below I'll walk through the specific issues I found and how the new design addresses each one.

Issue 1: Parents know the target's schema

This is the one that bothered me most. e.g. in Scorer._create_identifier():

# pyrit/score/scorer.py - _create_identifier()
target_info: Optional[Dict[str, Any]] = None
if prompt_target:
    target_id = prompt_target.get_identifier()
    target_info = {"class_name": target_id.class_name}
    if target_id.model_name:
        target_info["model_name"] = target_id.model_name
    if target_id.temperature is not None:
        target_info["temperature"] = target_id.temperature
    if target_id.top_p is not None:
        target_info["top_p"] = target_id.top_p

The scorer is reaching into the target's TargetIdentifier and cherry-picking fields it thinks are important (model_name, temperature, top_p). This is tight coupling, the scorer has to know the target's identity schema. If someone adds a field to TargetIdentifier that matters for behavior (say max_tokens or frequency_penalty), the scorer's identity won't reflect it unless someone also remembers to update this extraction logic.

And the extraction is lossy. The target has a full typed identifier with all its params; the scorer flattens it to a Dict[str, Any] with only the fields someone decided to copy over.

This isn't isolated to the scorer. The converter has the exact same coupling:

# pyrit/prompt_converter/prompt_converter.py -  _create_identifier()
if converter_target:
    target_id = converter_target.get_identifier()
    target_info = {
        "class_name": target_id.class_name,
        "model_name": target_id.model_name,
        "temperature": target_id.temperature,
        "top_p": target_id.top_p,
    }

Same pattern, the parent reaches into the child's schema and cherry-picks the same four fields. Actually worse than the scorer version because it doesn't even check for None values. Any component that holds a reference to another component and wants to include it in its identity has to repeat this extraction logic.

With ComponentConfig, the scorer doesn't look inside the target at all:

class SelfAskScaleScorer(FloatScaleScorer, Configurable):
    def _build_config(self) -> ComponentConfig:
        return self._create_config(
            params={
                self.CONFIG_KEY_SYSTEM_PROMPT: self._system_prompt,
                self.CONFIG_KEY_USER_PROMPT_TEMPLATE: "objective: {objective}\nresponse: {response}",
            },
            children={self.CONFIG_KEY_PROMPT_TARGET: self._prompt_target.get_config()},
        )

The target's entire config is captured as a child. The scorer's hash incorporates the target's hash, if the target's config changes (new model, different temperature, anything), the scorer's hash changes automatically. The scorer never needs to know what fields the target has.

Issue 2: Hashing is implicit and annotation-driven

The current system uses field metadata annotations to control what gets hashed:

# pyrit/identifiers/identifier.py
class _ExcludeFrom(Enum):
    HASH = "hash"
    STORAGE = "storage"

_EXPANSION_CATALOG: dict[_ExcludeFrom, set[_ExcludeFrom]] = {
    _ExcludeFrom.HASH: {_ExcludeFrom.HASH},
    _ExcludeFrom.STORAGE: {_ExcludeFrom.STORAGE, _ExcludeFrom.HASH},
}

# Then on each field:
class Identifier:
    class_description: str = field(metadata={_EXCLUDE: {_ExcludeFrom.STORAGE}})
    identifier_type: IdentifierType = field(metadata={_EXCLUDE: {_ExcludeFrom.STORAGE}})
    hash: str | None = field(default=None, compare=False, kw_only=True, metadata={_EXCLUDE: {_ExcludeFrom.HASH}})
    pyrit_version: str = field(default_factory=..., kw_only=True, metadata={_EXCLUDE: {_ExcludeFrom.HASH}})

And in the subclasses:

class ScorerIdentifier(Identifier):
    system_prompt_template: Optional[str] = field(default=None, metadata={_MAX_STORAGE_LENGTH: 100})
    user_prompt_template: Optional[str] = field(default=None, metadata={_MAX_STORAGE_LENGTH: 100})

Then hash computation iterates over all fields and checks each one:

def _compute_hash(self) -> str:
    hashable_dict = {
        f.name: getattr(self, f.name) for f in fields(self) if not _is_excluded_from_hash(f)
    }

This is clever, but it means that whether a field participates in hashing is determined by metadata on the field declaration, not by the structure of the data. When someone adds a new field to a subclass, they have to remember to annotate it correctly. There's no compiler error if you forget, the field just silently participates in (or is excluded from) the hash based on defaults.

The _EXPANSION_CATALOG adds another layer, STORAGE implicitly means HASH too, so _ExcludeFrom.STORAGE actually excludes from both storage and hashing. You have to understand this expansion system to reason about what any given field does.

With ComponentConfig, the separation is structural. There are only two buckets:

@dataclass(frozen=True)
class ComponentConfig:
    class_name: str                    # hashed (part of identity)
    class_module: str                  # hashed
    params: Dict[str, Any]            # hashed - this IS the behavioral config
    children: Dict[str, ...]          # hashed (via child hashes)
    hash: str                         # NOT hashed (it's the output)
    pyrit_version: str                # NOT hashed (metadata for storage)

If it's in params, it gets hashed. If it's not, it doesn't. There is no annotation system, no expansion catalog, no per-field metadata to get wrong. You can't accidentally put a field in the wrong bucket because the buckets are different structural locations, not metadata flags.

Issue 3: Registry metadata inherits identity machinery it never uses

ScenarioMetadata and InitializerMetadata extend Identifier:

@dataclass(frozen=True)
class ScenarioMetadata(Identifier):
    default_strategy: str
    all_strategies: tuple[str, ...]
    aggregate_strategies: tuple[str, ...]
    default_datasets: tuple[str, ...]
    max_dataset_size: Optional[int]

Through this inheritance, ScenarioMetadata gets: _compute_hash(), hash, unique_name, pyrit_version, to_dict(), from_dict(), identifier_type, _ExcludeFrom field processing, _MAX_STORAGE_LENGTH truncation. None of these are ever called on ScenarioMetadata instances - the registry builds them on-the-fly for CLI display and filtering, then discards them.

This isn't just dead code on the instance. It's coupling. If someone changes how Identifier._compute_hash() works, or how to_dict() serializes, ScenarioMetadata is implicitly affected. The hash is still computed in __post_init__ even though nobody reads it. If someone adds a field to Identifier, ScenarioMetadata inherits it. You can't tell from reading ScenarioMetadata which parts of Identifier it actually relies on.

The _build_metadata call sites also have to pass identity fields that are meaningless for class metadata:

# ScenarioRegistry._build_metadata — current
return ScenarioMetadata(
    identifier_type="class",             # always "class", why is this a per-instance field?
    class_name=scenario_class.__name__,
    class_module=scenario_class.__module__,
    class_description=description,
    default_strategy=...,
    all_strategies=...,
    ...
)

# InitializerRegistry._build_metadata - current
return InitializerMetadata(
    identifier_type="class",             # same - always "class"
    class_name=initializer_class.__name__,
    class_module=initializer_class.__module__,
    class_description=instance.description,
    display_name=instance.name,
    required_env_vars=tuple(instance.required_env_vars),
    execution_order=instance.execution_order,
)

Every call passes identifier_type="class". Every instance computes a hash that nobody reads. Every instance carries pyrit_version and unique_name that no consumer accesses. It's ceremony inherited from a base class designed for a different purpose.

With RegistryEntry, the base class has exactly what class registries need and nothing more:

@dataclass(frozen=True)
class RegistryEntry:
    class_name: str
    class_module: str
    class_description: str = ""

    @property
    def snake_class_name(self) -> str:
        return class_name_to_snake_case(self.class_name)

The metadata classes inherit only what they use:

# AFTER
@dataclass(frozen=True)
class ScenarioMetadata(RegistryEntry):
    default_strategy: str = ""
    all_strategies: tuple[str, ...] = ()
    aggregate_strategies: tuple[str, ...] = ()
    default_datasets: tuple[str, ...] = ()
    max_dataset_size: Optional[int] = None

@dataclass(frozen=True)
class InitializerMetadata(RegistryEntry):
    display_name: str = ""
    required_env_vars: tuple[str, ...] = ()
    execution_order: int = 100

And _build_metadata drops the dead fields:

# ScenarioRegistry._build_metadata — new
return ScenarioMetadata(
    class_name=scenario_class.__name__,   # no identifier_type
    class_module=scenario_class.__module__,
    class_description=description,
    default_strategy=...,
    all_strategies=...,
    ...
)

ScenarioMetadata(RegistryEntry) inherits 3 fields and 1 property. No hash computation, no serialization machinery, no version tracking, no exclusion annotations. Changes to ComponentConfig hashing or serialization can't affect registry metadata because they don't share a base class.

Because these metadata classes are ephemeral, created by _build_metadata for CLI display and filtering, never persisted to a database, the base class swap is a direct replacement. No dual-implementation period needed, unlike the Scorer/PromptTarget migration where both Identifiable and Configurable must coexist while consumers migrate.

The downstream consumers are unaffected:

  • CLI formatting (format_scenario_metadata, format_initializer_metadata) accesses snake_class_name, class_name, class_description, and domain-specific fields, all present on the new types.
  • _matches_filters is duck-typed via getattr and works with both Identifier and RegistryEntry.
  • BaseClassRegistry's MetadataT bound changes from Identifier to RegistryEntry.

Issue 4: Subclass-per-component is a typed schema for what's really an arbitrary param bag

We have ScorerIdentifier, TargetIdentifier, and ConverterIdentifier, each extending Identifier with their own fields. But look at what the subclasses actually add:

# ScorerIdentifier adds:
scorer_type: ScoreType
system_prompt_template: Optional[str]
user_prompt_template: Optional[str]
sub_identifier: Optional[List["ScorerIdentifier"]]
target_info: Optional[Dict[str, Any]]
score_aggregator: Optional[str]
scorer_specific_params: Optional[Dict[str, Any]]    # an easy way to escape the data structure

# TargetIdentifier adds:
endpoint: str
model_name: str
temperature: Optional[float]
top_p: Optional[float]
max_requests_per_minute: Optional[int]
target_specific_params: Optional[Dict[str, Any]]    # an easy way to escape the data structure

These fields are the behavioral parameters of each component, which model, what temperature, which prompt template. They're configuration written once at construction, frozen, and then serialized to JSON or hashed. Nobody writes identifier.scorer_type = "float_scale" during a running attack.

The tell is the escape hatches. Both subclasses have a *_specific_params: Optional[Dict[str, Any]] field, an untyped dict for "everything else that doesn't fit the schema." OpenAICompletionTarget uses it for max_tokens, frequency_penalty, presence_penalty, and n:

# pyrit/prompt_target/openai/openai_completion_target.py
def _build_identifier(self) -> TargetIdentifier:
    return self._create_identifier(
        temperature=self._temperature,
        top_p=self._top_p,
        target_specific_params={
            "max_tokens": self._max_tokens,
            "frequency_penalty": self._frequency_penalty,
            "presence_penalty": self._presence_penalty,
            "n": self._n,
        },
    )

Why do temperature and top_p get typed fields but max_tokens and frequency_penalty go into the untyped dict? There's no principled boundary, it's just which params someone happened to put on the base class first. The schema can't keep up with the variety of component params, so it falls back to a dict anyway.

And the main consumers interact with identifiers through the base-class surface:

# ScoreEntry stores it as JSON
entry.scorer_class_identifier = json.dumps(scorer_id.to_dict())

# Metrics lookup compares hashes
scorer_hash = scorer.get_identifier().hash

# ScenarioResult stores it as a dict
self.objective_target_identifier = target.get_identifier()

The one exception is ConsoleScorerPrinter, which accesses subclass fields directly, scorer_identifier.target_info, .score_aggregator, .sub_identifier. But even there, target_info is already Dict[str, Any] accessed via .get("model_name"), and scorer_specific_params is iterated as a generic dict. It's half-typed at best. With ComponentConfig this becomes config.params.get("score_aggregator") and config.children.get("sub_scorers", []), same access pattern, just through params and children instead of named fields.

Nobody branches on isinstance(identifier, ScorerIdentifier). The subclass-specific fields are written once at construction, serialized into a JSON blob in the DB, and queried by string key (scorer_class_identifier->>'scorer_type'). The type information is erased at the consumption boundary.

When the fields are arbitrary component params, the subclasses already need untyped dict escape hatches, and the single consumer that does typed access is already half-untyped, the subclass hierarchy is providing a typed schema for data that doesn't benefit from one.

With ComponentConfig, this is explicit:

# Scorer - params are whatever matters for this scorer
ComponentConfig.of(self, params={"scorer_type": ..., "system_prompt": ...}, children={"target": ...})

# Target - params are whatever matters for this target
ComponentConfig.of(self, params={"endpoint": ..., "model_name": ..., "temperature": ..., "max_tokens": ...})

# Converter - params are whatever matters for this converter
ComponentConfig.of(self, params={"converter_name": ..., "language": ...})

No escape hatch needed because the whole thing is a params dict. Components differ in what they put in params, not in what subclass they instantiate. Adding a new component type means implementing one method, no new Identifier subclass with its own from_dict() override.

Issue 5: The Identifiable generic adds ceremony for no benefit

Components currently must declare their identifier type:

class SelfAskScaleScorer(FloatScaleScorer, Identifiable[ScorerIdentifier]):
    def _build_identifier(self) -> ScorerIdentifier:
        return ScorerIdentifier(
            identifier_type="instance",
            class_name=self.__class__.__name__,
            class_module=self.__class__.__module__,
            class_description=" ".join(self.__class__.__doc__.split()) if self.__class__.__doc__ else "",
            ...
        )

The generic type parameter [ScorerIdentifier] looks like it gives type safety, but in practice callers almost always work with the result as a dict (to_dict()) or only access base-class fields (hash, class_name). The generic doesn't prevent misuse, it just adds boilerplate.

And every _build_identifier has to manually pass class_name=self.__class__.__name__, class_module=self.__class__.__module__, class_description=..., identifier_type="instance", the same 4 lines repeated in every implementation.

With Configurable, the boilerplate is gone:

class SelfAskScaleScorer(FloatScaleScorer, Configurable):
   def _build_config(self) -> ComponentConfig:
        return self._create_config(
            params={
                self.CONFIG_KEY_SYSTEM_PROMPT: self._system_prompt,
                self.CONFIG_KEY_USER_PROMPT_TEMPLATE: "objective: {objective}\nresponse: {response}",
            },
            children={self.CONFIG_KEY_PROMPT_TARGET: self._prompt_target.get_config()},
        )

   # base class create config
   def _create_config(
        self,
        *,
        params: Optional[Dict[str, Any]] = None,
        children: Optional[Dict[str, Union[ComponentConfig, List[ComponentConfig]]]] = None,
    ) -> ComponentConfig:
        all_params: Dict[str, Any] = {self.CONFIG_KEY_SCORER_TYPE: self.scorer_type}
        if params:
            all_params.update(params)

        return ComponentConfig.of(
             self,  # class_name and class_module extracted automatically
             params=all_params, 
             children=children
        )

ComponentConfig.of(self) extracts class_name and class_module from the object. No generic type parameter, no manual boilerplate fields.


The new design

Two types, each with one job:

ComponentConfig - instance identity. A frozen dataclass with (class_name, class_module, params, children, hash). Replaces ScorerIdentifier, TargetIdentifier, ConverterIdentifier, and Identifiable[T].

RegistryEntry - class metadata. A frozen dataclass with (class_name, class_module, class_description) and a snake_class_name property. Replaces Identifier as the base for ScenarioMetadata and InitializerMetadata.

These two types share no base class. They're connected only by duck typing - _matches_filters() uses getattr and works with both. Changes to one side can't break the other.

Summary of how each issue is resolved

Issue Current New
Scorer/converter knows target's schema Both manually extract model_name, temperature, top_p from TargetIdentifier Child config carried as opaque ComponentConfig, hash flows automatically
Hash behavior hidden in field metadata _ExcludeFrom enum, _EXPANSION_CATALOG, per-field metadata={_EXCLUDE: ...} Structural: params gets hashed, everything else doesn't
Registry metadata inherits unused machinery ScenarioMetadata(Identifier) inherits hash, serialization, exclusion logic, identifier_type="class" on every instance ScenarioMetadata(RegistryEntry) inherits 3 fields and 1 property; direct swap since metadata is ephemeral
Subclass per component type ScorerIdentifier, TargetIdentifier, ConverterIdentifier with typed schemas for arbitrary params (plus *_specific_params dict escape hatches) One ComponentConfig type, components differ in params
Generic ABC boilerplate Identifiable[ScorerIdentifier], manual class_name/class_module/identifier_type Configurable ABC, ComponentConfig.of(self) extracts class info

Other things worth noting

Backward-compatible schema evolution. Adding an optional param with None default doesn't change existing hashes because None values are excluded:

config_v1 = ComponentConfig.of(self, params={"threshold": 0.8})
config_v2 = ComponentConfig.of(self, params={"threshold": 0.8, "new_param": None})
assert config_v1.hash == config_v2.hash

DB backward compatibility. to_dict() inlines params at top level, so existing queries like scorer_class_identifier->>'scorer_type' keep working. from_dict() handles legacy __type__/__module__ keys.

Open/Closed for new components. Adding a new component type means implementing _build_config() and putting params in a dict. No new identifier subclass, no new from_dict() override.

Architecture

The design separates instance identity from class metadata, with no shared base class between them:

INSTANCE IDENTITY                           CLASS METADATA
(DB storage, hashing, metrics)              (registries, CLI display)

ComponentConfig                             RegistryEntry
├── class_name                              ├── class_name
├── class_module                            ├── class_module
├── params: dict                            ├── class_description
├── children: dict                          └── snake_class_name (property)
├── hash (content-addressed)                    │
├── pyrit_version                               ├── ScenarioMetadata
├── snake_class_name (property)                 │   ├── default_strategy
├── unique_name (property)                      │   ├── all_strategies
└── short_hash (property)                       │   ├── aggregate_strategies
                                                │   ├── default_datasets
Configurable (ABC)                              │   └── max_dataset_size
├── _build_config()                             │
└── get_config() (cached)                       └── InitializerMetadata
                                                    ├── display_name
                                                    ├── required_env_vars
                                                    └── execution_order

Why two unrelated types? Because they have fundamentally different lifecycles and consumers:

Concern ComponentConfig RegistryEntry
Created when Component is instantiated _build_metadata called for CLI/listing
Persisted to DB Yes (Score, MessagePiece, ScenarioResult) No - ephemeral
Needs hashing Yes - content-addressed identity No
Needs serialization Yes - to_dict()/from_dict() for DB roundtrip No
Needs version tracking Yes - pyrit_version for schema evolution No
Connected via _matches_filters uses getattr on both (duck-typed) Same

Migration plan

Incremental, no big-bang:

  1. Add new types alongside existing - ComponentConfig, Configurable, RegistryEntry added. Zero breakage, both systems coexist.
  2. Add ComponentConfig.normalize() - classmethod that accepts ComponentConfig | dict and returns ComponentConfig, mirroring the existing Identifier.normalize() pattern. Needed because MessagePiece uses .normalize() to handle both typed identifiers and legacy dicts from the DB.
  3. Dual-implement on components - scorers, targets, converters implement both Configurable and Identifiable temporarily, verify hashes match.
  4. Swap class registry metadata - ScenarioMetadata(Identifier)ScenarioMetadata(RegistryEntry), InitializerMetadata(Identifier) -> InitializerMetadata(RegistryEntry), BaseClassRegistry's MetadataT bound: Identifier -> RegistryEntry. This is a direct replacement (no dual-implementation) because metadata objects are ephemeral, created by _build_metadata for CLI display, never persisted. The only code change in _build_metadata call sites is removing identifier_type="class".
  5. Migrate MessagePiece - this is the largest consumer surface. Three typed identifier fields (converter_identifiers: List[ConverterIdentifier], prompt_target_identifier: Optional[TargetIdentifier], scorer_identifier: Optional[ScorerIdentifier]) change to ComponentConfig. PromptMemoryEntry.get_message_piece() switches from TargetIdentifier.from_dict() / ConverterIdentifier.from_dict() to ComponentConfig.from_dict().
  6. Migrate other consumers - ScenarioResult, Score, ConsoleScorerPrinter, ScorerRegistry, TargetRegistry, metrics IO.
  7. Remove old types - delete Identifier hierarchy, Identifiable, exclusion machinery.

Deliberate simplification: The current _MAX_STORAGE_LENGTH metadata truncates long prompt templates in to_dict() for storage display. ComponentConfig doesn't replicate this - if a UI or printer wants to truncate long values, that's a display concern, not an identity concern. The hash already uses full values in both designs.

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.

1 participant