FEAT: Add TargetRegistry.get_by_tag_query for TagQuery-based lookup

Author: Victor Valbuena

pyrit/registry/object_registries/target_registry.py

@@ -12,9 +12,11 @@
 import logging
 from typing import TYPE_CHECKING, Optional, Union
 
+from pyrit.registry.object_registries.base_instance_registry import RegistryEntry
 from pyrit.registry.object_registries.retrievable_instance_registry import (
     RetrievableInstanceRegistry,
 )
+from pyrit.registry.tag_query import TagQuery
 
 if TYPE_CHECKING:
     from pyrit.prompt_target import PromptTarget
@@ -74,3 +76,30 @@ def get_instance_by_name(self, name: str) -> Optional[PromptTarget]:
             The target instance, or None if not found.
         """
         return self.get(name)
+
+    def get_by_tag_query(self, *, query: TagQuery) -> list[RegistryEntry[PromptTarget]]:
+        """
+        Get all entries whose tag keys satisfy ``query``.
+
+        ``TagQuery`` operates on a tag set, so this method matches against
+        ``entry.tags.keys()`` and ignores tag values. For value-aware
+        single-tag lookups use ``get_by_tag(*, tag, value)`` on the base
+        class.
+
+        Composite queries compose with ``&`` and ``|`` operators, e.g.
+        ``TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn")``.
+
+        Args:
+            query: The tag predicate to evaluate against each entry.
+
+        Returns:
+            List of matching ``RegistryEntry`` objects sorted by registry name.
+        """
+        results: list[RegistryEntry[PromptTarget]] = []
+        # Note: this erases insertion order, but respects the base_instance_registry pattern
+        # (get_by_tag).
+        for name in sorted(self._registry_items.keys()):
+            entry = self._registry_items[name]
+            if query.matches(set(entry.tags.keys())):
+                results.append(entry)
+        return results

tests/unit/registry/test_target_registry.py

@@ -230,3 +230,64 @@ def test_list_metadata_filter_by_class_name(self):
         assert len(mock_metadata) == 2
         for m in mock_metadata:
             assert m.class_name == "MockPromptTarget"
+
+
+@pytest.mark.usefixtures("patch_central_database")
+class TestTargetRegistryGetByTagQuery:
+    """Tests for ``TargetRegistry.get_by_tag_query`` (TagQuery-aware tag lookup)."""
+
+    def setup_method(self):
+        """Reset and populate a fresh registry for each test."""
+        TargetRegistry.reset_instance()
+        self.registry = TargetRegistry.get_registry_singleton()
+
+        self.registry.register_instance(MockPromptTarget(), name="adv_single", tags=["adversarial", "singleturn"])
+        self.registry.register_instance(MockPromptTarget(), name="adv_multi", tags=["adversarial", "multiturn"])
+        self.registry.register_instance(MockPromptChatTarget(), name="scorer_only", tags=["scorer"])
+        self.registry.register_instance(MockPromptTarget(), name="untagged")
+
+    def teardown_method(self):
+        """Reset the singleton after each test."""
+        TargetRegistry.reset_instance()
+
+    def test_get_by_tag_query_returns_matching(self):
+        """A leaf ``TagQuery.all`` returns every entry whose tag set contains the required tag."""
+        from pyrit.registry.tag_query import TagQuery
+
+        results = self.registry.get_by_tag_query(query=TagQuery.all("adversarial"))
+
+        names = [entry.name for entry in results]
+        assert names == ["adv_multi", "adv_single"]
+
+    def test_get_by_tag_query_empty(self):
+        """A query that matches no entries returns an empty list (not raise)."""
+        from pyrit.registry.tag_query import TagQuery
+
+        results = self.registry.get_by_tag_query(query=TagQuery.all("nonexistent_tag"))
+        assert results == []
+
+    def test_get_by_tag_query_composite_and_or(self):
+        """Composite queries via ``&`` / ``|`` evaluate as expected."""
+        from pyrit.registry.tag_query import TagQuery
+
+        query = TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn")
+        results = self.registry.get_by_tag_query(query=query)
+
+        names = [entry.name for entry in results]
+        assert names == ["adv_multi", "adv_single"]
+
+        narrower = TagQuery.all("adversarial") & TagQuery.any_of("singleturn")
+        narrow_names = [entry.name for entry in self.registry.get_by_tag_query(query=narrower)]
+        assert narrow_names == ["adv_single"]
+
+    def test_get_by_tag_query_matches_keys_not_values(self):
+        """``TagQuery`` evaluates against tag keys; tag values are ignored by this method."""
+        from pyrit.registry.tag_query import TagQuery
+
+        self.registry.add_tags(name="adv_single", tags={"priority": "high"})
+
+        priority_matches = self.registry.get_by_tag_query(query=TagQuery.all("priority"))
+        assert [entry.name for entry in priority_matches] == ["adv_single"]
+
+        value_lookup = self.registry.get_by_tag_query(query=TagQuery.all("high"))
+        assert value_lookup == []