Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion pyrit/prompt_target/azure_ml_chat_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,10 @@ class AzureMLChatTarget(PromptChatTarget):
api_key_environment_variable: str = "AZURE_ML_KEY"

_DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
supports_multi_message_pieces=True, supports_editable_history=True, supports_multi_turn=True
supports_multi_message_pieces=True,
supports_editable_history=True,
supports_multi_turn=True,
supports_system_prompt=True,
)

def __init__(
Expand Down
4 changes: 3 additions & 1 deletion pyrit/prompt_target/common/prompt_chat_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ class PromptChatTarget(PromptTarget):
"""

_DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
supports_multi_turn=True, supports_multi_message_pieces=True
supports_multi_turn=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
)

def __init__(
Expand Down
5 changes: 5 additions & 0 deletions pyrit/prompt_target/common/target_capabilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ class attribute. Users can override individual capabilities per instance
# multi-turn interactions and that the attack history is not immutable once set.
supports_editable_history: bool = False

# Whether the target natively supports system prompts.
supports_system_prompt: bool = False

# The input modalities supported by the target (e.g., "text", "image").
input_modalities: frozenset[frozenset[PromptDataType]] = frozenset({frozenset(["text"])})

Expand Down Expand Up @@ -76,6 +79,7 @@ def get_known_capabilities(underlying_model: str) -> "Optional[TargetCapabilitie
_GPT_4O = TargetCapabilities(
supports_multi_turn=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
supports_json_output=True,
input_modalities=_TEXT_IMAGE_INPUT,
output_modalities=_TEXT_OUTPUT,
Expand All @@ -84,6 +88,7 @@ def get_known_capabilities(underlying_model: str) -> "Optional[TargetCapabilitie
_GPT_5 = TargetCapabilities(
supports_multi_turn=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
supports_json_schema=True,
supports_json_output=True,
input_modalities=_TEXT_IMAGE_INPUT,
Expand Down
4 changes: 3 additions & 1 deletion pyrit/prompt_target/hugging_face/hugging_face_chat_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,9 @@ class HuggingFaceChatTarget(PromptChatTarget):
"""

_DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
supports_multi_turn=True, supports_editable_history=True
supports_multi_turn=True,
supports_editable_history=True,
supports_system_prompt=True,
)

# Class-level cache for model and tokenizer
Expand Down
1 change: 1 addition & 0 deletions pyrit/prompt_target/openai/openai_chat_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ class OpenAIChatTarget(OpenAITarget, PromptChatTarget):
supports_multi_turn=True,
supports_json_output=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
)

def __init__(
Expand Down
1 change: 1 addition & 0 deletions pyrit/prompt_target/openai/openai_realtime_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ class RealtimeTarget(OpenAITarget, PromptChatTarget):
_DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
supports_multi_turn=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
input_modalities=frozenset(
{
frozenset(["text"]),
Expand Down
1 change: 1 addition & 0 deletions pyrit/prompt_target/openai/openai_response_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ class OpenAIResponseTarget(OpenAITarget, PromptChatTarget):
supports_multi_turn=True,
supports_json_output=True,
supports_multi_message_pieces=True,
supports_system_prompt=True,
input_modalities=frozenset(
{
frozenset(["text"]),
Expand Down
35 changes: 35 additions & 0 deletions tests/unit/target/test_target_capabilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,7 @@ def test_hugging_face_chat_target_capabilities(self):
caps = HuggingFaceChatTarget._DEFAULT_CAPABILITIES
assert caps.supports_editable_history is True
assert caps.supports_multi_turn is True
assert caps.supports_system_prompt is True

def test_azure_ml_chat_target_capabilities(self):
from pyrit.prompt_target import AzureMLChatTarget
Expand All @@ -253,6 +254,31 @@ def test_azure_ml_chat_target_capabilities(self):
)
assert target.capabilities.supports_editable_history is True
assert target.capabilities.supports_multi_message_pieces is True
assert target.capabilities.supports_system_prompt is True

@patch.dict("os.environ", _CLEAN_UNDERLYING_MODEL_ENV)
def test_prompt_chat_targets_support_system_prompt(self):
from pyrit.prompt_target import OpenAIChatTarget, OpenAIResponseTarget, RealtimeTarget

openai_chat_target = OpenAIChatTarget(
model_name="test-model",
endpoint="https://mock.azure.com/",
api_key="mock-api-key",
)
openai_response_target = OpenAIResponseTarget(
model_name="o1",
endpoint="https://mock.azure.com/",
api_key="mock-api-key",
)
realtime_target = RealtimeTarget(
model_name="gpt-4o-realtime",
endpoint="https://mock.azure.com/",
api_key="mock-api-key",
)

assert openai_chat_target.capabilities.supports_system_prompt is True
assert openai_response_target.capabilities.supports_system_prompt is True
assert realtime_target.capabilities.supports_system_prompt is True

def test_custom_capabilities_override_modalities(self):
from pyrit.prompt_target import OpenAIChatTarget, TargetCapabilities
Expand Down Expand Up @@ -415,3 +441,12 @@ def test_recognized_model_overrides_class_default(self):
cls = self._make_target_class(default_caps=minimal_caps)
result = cls.get_default_capabilities("tts")
assert result.output_modalities == frozenset({frozenset(["audio_path"])})

def test_prompt_chat_target_preserves_system_prompt_for_recognized_model(self):
from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget

result = PromptChatTarget.get_default_capabilities("gpt-4o")

assert result.supports_multi_turn is True
assert result.supports_multi_message_pieces is True
assert result.supports_system_prompt is True
Loading