diff --git a/pyrit/prompt_target/azure_ml_chat_target.py b/pyrit/prompt_target/azure_ml_chat_target.py index cc83332027..4795d42d0d 100644 --- a/pyrit/prompt_target/azure_ml_chat_target.py +++ b/pyrit/prompt_target/azure_ml_chat_target.py @@ -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__( diff --git a/pyrit/prompt_target/common/prompt_chat_target.py b/pyrit/prompt_target/common/prompt_chat_target.py index af12a11a1e..6b77ac03b1 100644 --- a/pyrit/prompt_target/common/prompt_chat_target.py +++ b/pyrit/prompt_target/common/prompt_chat_target.py @@ -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__( diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 5b1b596120..e6ced6a1a2 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -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"])}) @@ -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, @@ -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, diff --git a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py index 8eeb0bfc6a..13152af775 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py @@ -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 diff --git a/pyrit/prompt_target/openai/openai_chat_target.py b/pyrit/prompt_target/openai/openai_chat_target.py index dcc7409934..ba0abf7cfa 100644 --- a/pyrit/prompt_target/openai/openai_chat_target.py +++ b/pyrit/prompt_target/openai/openai_chat_target.py @@ -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__( diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index 22edf20391..ffc8d7a2ca 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -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"]), diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index be720fa5c9..b41b40a30f 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -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"]), diff --git a/tests/unit/target/test_target_capabilities.py b/tests/unit/target/test_target_capabilities.py index f64d7d3203..48b62297c3 100644 --- a/tests/unit/target/test_target_capabilities.py +++ b/tests/unit/target/test_target_capabilities.py @@ -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 @@ -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 @@ -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