-
Notifications
You must be signed in to change notification settings - Fork 649
FEAT: Target Registry and AIRT Targets Initializer #1320
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
jsong468
wants to merge
1
commit into
Azure:main
Choose a base branch
from
jsong468:target_registry
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+682
−2
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,145 @@ | ||
| # Copyright (c) Microsoft Corporation. | ||
| # Licensed under the MIT license. | ||
|
|
||
| """ | ||
| Target registry for discovering and managing PyRIT prompt targets. | ||
|
|
||
| Targets are registered explicitly via initializers as pre-configured instances. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import hashlib | ||
| import json | ||
| import logging | ||
| from dataclasses import dataclass | ||
| from typing import TYPE_CHECKING, Any, Dict, Optional | ||
|
|
||
| from pyrit.registry.base import RegistryItemMetadata | ||
| from pyrit.registry.instance_registries.base_instance_registry import ( | ||
| BaseInstanceRegistry, | ||
| ) | ||
| from pyrit.registry.name_utils import class_name_to_registry_name | ||
|
|
||
| if TYPE_CHECKING: | ||
| from pyrit.prompt_target import PromptTarget | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class TargetMetadata(RegistryItemMetadata): | ||
| """ | ||
| Metadata describing a registered target instance. | ||
|
|
||
| Unlike ScenarioMetadata/InitializerMetadata which describe classes, | ||
| TargetMetadata describes an already-instantiated prompt target. | ||
|
|
||
| Use get() to retrieve the actual target instance. | ||
| """ | ||
|
|
||
| target_identifier: Dict[str, Any] | ||
|
|
||
|
|
||
| class TargetRegistry(BaseInstanceRegistry["PromptTarget", TargetMetadata]): | ||
| """ | ||
| Registry for managing available prompt target instances. | ||
|
|
||
| This registry stores pre-configured PromptTarget instances (not classes). | ||
| Targets are registered explicitly via initializers after being instantiated | ||
| with their required parameters (e.g., endpoint, API keys). | ||
|
|
||
| Targets are identified by their snake_case name derived from the class name, | ||
| or a custom name provided during registration. | ||
| """ | ||
|
|
||
| @classmethod | ||
| def get_registry_singleton(cls) -> "TargetRegistry": | ||
| """ | ||
| Get the singleton instance of the TargetRegistry. | ||
|
|
||
| Returns: | ||
| The singleton TargetRegistry instance. | ||
| """ | ||
| return super().get_registry_singleton() # type: ignore[return-value] | ||
|
|
||
| def register_instance( | ||
| self, | ||
| target: "PromptTarget", | ||
| *, | ||
| name: Optional[str] = None, | ||
| ) -> None: | ||
| """ | ||
| Register a target instance. | ||
|
|
||
| Note: Unlike ScenarioRegistry and InitializerRegistry which register classes, | ||
| TargetRegistry registers pre-configured instances. | ||
|
|
||
| Args: | ||
| target: The pre-configured target instance (not a class). | ||
| name: Optional custom registry name. If not provided, | ||
| derived from class name with identifier hash appended | ||
| (e.g., OpenAIChatTarget -> openai_chat_abc123). | ||
| """ | ||
| if name is None: | ||
| base_name = class_name_to_registry_name(target.__class__.__name__, suffix="Target") | ||
| # Append identifier hash for uniqueness | ||
| identifier_hash = self._compute_identifier_hash(target)[:8] | ||
| name = f"{base_name}_{identifier_hash}" | ||
|
|
||
| self.register(target, name=name) | ||
| logger.debug(f"Registered target instance: {name} ({target.__class__.__name__})") | ||
|
|
||
| def get_instance_by_name(self, name: str) -> Optional["PromptTarget"]: | ||
| """ | ||
| Get a registered target instance by name. | ||
|
|
||
| Note: This returns an already-instantiated target, not a class. | ||
|
|
||
| Args: | ||
| name: The registry name of the target. | ||
|
|
||
| Returns: | ||
| The target instance, or None if not found. | ||
| """ | ||
| return self.get(name) | ||
|
|
||
| def _build_metadata(self, name: str, instance: "PromptTarget") -> TargetMetadata: | ||
| """ | ||
| Build metadata for a target instance. | ||
|
|
||
| Args: | ||
| name: The registry name of the target. | ||
| instance: The target instance. | ||
|
|
||
| Returns: | ||
| TargetMetadata describing the target. | ||
| """ | ||
| # Get description from docstring | ||
| doc = instance.__class__.__doc__ or "" | ||
| description = " ".join(doc.split()) if doc else "No description available" | ||
|
|
||
| # Get identifier from the target | ||
| target_identifier = instance.get_identifier() | ||
|
|
||
| return TargetMetadata( | ||
| name=name, | ||
| class_name=instance.__class__.__name__, | ||
| description=description, | ||
| target_identifier=target_identifier, | ||
| ) | ||
|
|
||
| @staticmethod | ||
| def _compute_identifier_hash(target: "PromptTarget") -> str: | ||
| """ | ||
| Compute a hash from the target's identifier for unique naming. | ||
|
|
||
| Args: | ||
| target: The target instance. | ||
|
|
||
| Returns: | ||
| A hex string hash of the identifier. | ||
| """ | ||
| identifier = target.get_identifier() | ||
| identifier_str = json.dumps(identifier, sort_keys=True) | ||
| return hashlib.sha256(identifier_str.encode()).hexdigest() | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like the idea of an identifier hash, but I think if we want to add it, it should be general. I also think we should have a minor refactor while it's still early.
Identifierclass which is inmodelsmoduleI can create a PR for the Identifiers refactor, and there is a story in the backlog to do ScorerIdentifier. But I think as part of this we should move TargetIdentifier to a dataclass also
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For you, I might make a separate (prereq) PR that just converts the TargetIdentifier to the dataclass and we can get that in first as I work on the other piece.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
#1323