Skip to content

Commit 8b812dd

Browse files
authored
MAINT: Allow custom Likert system prompt and scale (microsoft#1514)
1 parent 4e69237 commit 8b812dd

2 files changed

Lines changed: 303 additions & 17 deletions

File tree

pyrit/score/float_scale/self_ask_likert_scorer.py

Lines changed: 104 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -166,36 +166,63 @@ def __init__(
166166
self,
167167
*,
168168
chat_target: PromptChatTarget,
169-
likert_scale: LikertScalePaths,
169+
likert_scale: Optional[LikertScalePaths] = None,
170+
custom_likert_path: Optional[Path] = None,
171+
custom_system_prompt_path: Optional[Path] = None,
170172
validator: Optional[ScorerPromptValidator] = None,
171173
) -> None:
172174
"""
173175
Initialize the SelfAskLikertScorer.
174176
175177
Args:
176178
chat_target (PromptChatTarget): The chat target to use for scoring.
177-
likert_scale (LikertScalePaths): The Likert scale configuration to use for scoring.
179+
likert_scale (Optional[LikertScalePaths]): The Likert scale configuration to use for scoring.
180+
custom_likert_path (Optional[Path]): Path to a custom YAML file containing the Likert scale definition.
181+
This allows users to use their own Likert scales without modifying the code, as long as
182+
the YAML file follows the expected format. Only one of `likert_scale` or `custom_likert_path`
183+
should be provided. Defaults to None.
184+
custom_system_prompt_path (Optional[Path]): Path to a custom system prompt file. This allows users to
185+
provide their own system prompt without modifying the code. Defaults to None.
178186
validator (Optional[ScorerPromptValidator]): Custom validator for the scorer. Defaults to None.
187+
188+
Raises:
189+
ValueError: If both `likert_scale` and `custom_likert_path` are provided, if neither is provided,
190+
or if the provided Likert scale or system prompt YAML file is improperly formatted.
179191
"""
180192
super().__init__(validator=validator or self._DEFAULT_VALIDATOR)
181193

182194
self._prompt_target = chat_target
183195
self._likert_scale = likert_scale
184196

185-
# Auto-set evaluation file mapping from the LikertScalePaths enum
186-
if likert_scale.evaluation_files is not None:
187-
from pyrit.score.scorer_evaluation.scorer_evaluator import (
188-
ScorerEvalDatasetFiles,
189-
)
197+
if likert_scale is not None and custom_likert_path is not None:
198+
raise ValueError("Only one of 'likert_scale' or 'custom_likert_path' should be provided, not both.")
199+
if likert_scale is None and custom_likert_path is None:
200+
raise ValueError("One of 'likert_scale' or 'custom_likert_path' must be provided.")
190201

191-
eval_files = likert_scale.evaluation_files
192-
self.evaluation_file_mapping = ScorerEvalDatasetFiles(
193-
human_labeled_datasets_files=eval_files.human_labeled_datasets_files,
194-
result_file=eval_files.result_file,
195-
harm_category=eval_files.harm_category,
196-
)
202+
self._scoring_instructions_template: Optional[SeedPrompt] = (
203+
None # Will be set in _set_likert_scale_system_prompt
204+
)
205+
if custom_system_prompt_path is not None:
206+
self._validate_custom_system_prompt_path(custom_system_prompt_path)
207+
self._scoring_instructions_template = SeedPrompt.from_yaml_file(custom_system_prompt_path)
208+
if likert_scale is not None:
209+
# Auto-set evaluation file mapping from the LikertScalePaths enum
210+
if likert_scale.evaluation_files is not None:
211+
from pyrit.score.scorer_evaluation.scorer_evaluator import (
212+
ScorerEvalDatasetFiles,
213+
)
214+
215+
eval_files = likert_scale.evaluation_files
216+
self.evaluation_file_mapping = ScorerEvalDatasetFiles(
217+
human_labeled_datasets_files=eval_files.human_labeled_datasets_files,
218+
result_file=eval_files.result_file,
219+
harm_category=eval_files.harm_category,
220+
)
197221

198-
self._set_likert_scale_system_prompt(likert_scale_path=likert_scale.path)
222+
self._set_likert_scale_system_prompt(likert_scale_path=likert_scale.path)
223+
elif custom_likert_path is not None:
224+
self._validate_custom_likert_path(custom_likert_path)
225+
self._set_likert_scale_system_prompt(likert_scale_path=custom_likert_path)
199226

200227
def _build_identifier(self) -> ComponentIdentifier:
201228
"""
@@ -268,9 +295,12 @@ def _set_likert_scale_system_prompt(self, likert_scale_path: Path) -> None:
268295
f"but only a single unique value was found: {self._max_scale_value}."
269296
)
270297

271-
self._scoring_instructions_template = SeedPrompt.from_yaml_file(
272-
SCORER_LIKERT_PATH / "likert_system_prompt.yaml"
273-
)
298+
# Only load the default system prompt template if a custom one wasn't already
299+
# set via custom_system_prompt_path in __init__.
300+
if self._scoring_instructions_template is None:
301+
self._scoring_instructions_template = SeedPrompt.from_yaml_file(
302+
SCORER_LIKERT_PATH / "likert_system_prompt.yaml"
303+
)
274304

275305
self._system_prompt = self._scoring_instructions_template.render_template_value(
276306
likert_scale=likert_scale_str,
@@ -337,6 +367,63 @@ def _likert_scale_description_to_string(self, descriptions: list[dict[str, str]]
337367

338368
return likert_scale_description
339369

370+
@staticmethod
371+
def _validate_custom_system_prompt_path(custom_system_prompt_path: Path) -> None:
372+
"""
373+
Validate the custom system prompt path.
374+
375+
Checks that the file exists, has a YAML extension, and contains the required
376+
template parameters (category, likert_scale, min_scale_value, max_scale_value)
377+
that the Likert scorer needs to render the system prompt.
378+
379+
Args:
380+
custom_system_prompt_path (Path): Path to the custom system prompt YAML file.
381+
382+
Raises:
383+
FileNotFoundError: If the file does not exist.
384+
ValueError: If the file is not a YAML file or is missing required template parameters.
385+
"""
386+
if not custom_system_prompt_path.exists():
387+
raise FileNotFoundError(f"Custom system prompt file not found: '{custom_system_prompt_path}'")
388+
if custom_system_prompt_path.suffix not in (".yaml", ".yml"):
389+
raise ValueError(
390+
f"Custom system prompt file must be a YAML file (.yaml or .yml), "
391+
f"got '{custom_system_prompt_path.suffix}'."
392+
)
393+
394+
# Validate the template contains all required parameters used by the Likert scorer.
395+
SeedPrompt.from_yaml_with_required_parameters(
396+
template_path=custom_system_prompt_path,
397+
required_parameters=["category", "likert_scale", "min_scale_value", "max_scale_value"],
398+
error_message=(
399+
"Custom system prompt YAML must define parameters: "
400+
"category, likert_scale, min_scale_value, max_scale_value"
401+
),
402+
)
403+
404+
@staticmethod
405+
def _validate_custom_likert_path(custom_likert_path: Path) -> None:
406+
"""
407+
Validate the custom Likert scale path.
408+
409+
Performs basic path checks (existence and YAML extension). Deeper content
410+
validation (category, scale_descriptions structure, score values) is handled
411+
by ``_set_likert_scale_system_prompt`` when the file is actually parsed.
412+
413+
Args:
414+
custom_likert_path (Path): Path to the custom Likert scale YAML file.
415+
416+
Raises:
417+
FileNotFoundError: If the file does not exist.
418+
ValueError: If the file is not a YAML file.
419+
"""
420+
if not custom_likert_path.exists():
421+
raise FileNotFoundError(f"Custom Likert scale file not found: '{custom_likert_path}'")
422+
if custom_likert_path.suffix not in (".yaml", ".yml"):
423+
raise ValueError(
424+
f"Custom Likert scale file must be a YAML file (.yaml or .yml), got '{custom_likert_path.suffix}'."
425+
)
426+
340427
async def _score_piece_async(self, message_piece: MessagePiece, *, objective: Optional[str] = None) -> list[Score]:
341428
"""
342429
Score the given message_piece using "self-ask" for the chat target.

tests/unit/score/test_self_ask_likert.py

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,3 +403,202 @@ def test_likert_scale_missing_score_value_key_rejected(tmp_path: Path):
403403
chat_target=chat_target,
404404
likert_scale=LikertScalePaths.CYBER_SCALE,
405405
)
406+
407+
408+
# ---------------------------------------------------------------------------
409+
# custom_likert_path and custom_system_prompt_path tests
410+
# ---------------------------------------------------------------------------
411+
412+
413+
def _make_custom_system_prompt_yaml(tmp_path: Path, *, include_all_params: bool = True) -> Path:
414+
"""Create a custom system prompt YAML file for testing."""
415+
params = ["category", "likert_scale", "min_scale_value", "max_scale_value"] if include_all_params else ["category"]
416+
prompt_data = {
417+
"name": "custom test prompt",
418+
"description": "test",
419+
"parameters": params,
420+
"data_type": "text",
421+
"value": "Custom prompt for {{category}} with scale {{likert_scale}} "
422+
"from {{min_scale_value}} to {{max_scale_value}}."
423+
if include_all_params
424+
else "Only {{category}}.",
425+
}
426+
yaml_file = tmp_path / "custom_system_prompt.yaml"
427+
yaml_file.write_text(yaml.safe_dump(prompt_data), encoding="utf-8")
428+
return yaml_file
429+
430+
431+
def test_custom_likert_path_creates_scorer(tmp_path: Path):
432+
"""Verify that passing custom_likert_path (instead of a LikertScalePaths enum) works."""
433+
memory = MagicMock(MemoryInterface)
434+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
435+
chat_target = MagicMock()
436+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
437+
438+
custom_path = _make_custom_scale_yaml(tmp_path, category="custom_cat", min_val=0, max_val=3)
439+
scorer = SelfAskLikertScorer(chat_target=chat_target, custom_likert_path=custom_path)
440+
441+
assert scorer._min_scale_value == 0
442+
assert scorer._max_scale_value == 3
443+
assert scorer._score_category == "custom_cat"
444+
445+
446+
def test_custom_likert_path_file_not_found():
447+
"""Verify that a non-existent custom_likert_path raises FileNotFoundError."""
448+
memory = MagicMock(MemoryInterface)
449+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
450+
chat_target = MagicMock()
451+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
452+
453+
with pytest.raises(FileNotFoundError, match="Custom Likert scale file not found"):
454+
SelfAskLikertScorer(chat_target=chat_target, custom_likert_path=Path("/does/not/exist.yaml"))
455+
456+
457+
def test_custom_likert_path_non_yaml_rejected(tmp_path: Path):
458+
"""Verify that a non-YAML custom_likert_path raises ValueError."""
459+
bad_file = tmp_path / "scale.txt"
460+
bad_file.write_text("not yaml", encoding="utf-8")
461+
462+
memory = MagicMock(MemoryInterface)
463+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
464+
chat_target = MagicMock()
465+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
466+
467+
with pytest.raises(ValueError, match="must be a YAML file"):
468+
SelfAskLikertScorer(chat_target=chat_target, custom_likert_path=bad_file)
469+
470+
471+
def test_custom_system_prompt_non_yaml_rejected(tmp_path: Path):
472+
"""Verify that a non-YAML custom_system_prompt_path raises ValueError."""
473+
bad_file = tmp_path / "prompt.txt"
474+
bad_file.write_text("not yaml", encoding="utf-8")
475+
476+
memory = MagicMock(MemoryInterface)
477+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
478+
chat_target = MagicMock()
479+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
480+
481+
with pytest.raises(ValueError, match="must be a YAML file"):
482+
SelfAskLikertScorer(
483+
chat_target=chat_target,
484+
likert_scale=LikertScalePaths.CYBER_SCALE,
485+
custom_system_prompt_path=bad_file,
486+
)
487+
488+
489+
def test_custom_system_prompt_path_used_in_system_prompt(tmp_path: Path):
490+
"""Verify that a custom system prompt template is rendered instead of the default."""
491+
memory = MagicMock(MemoryInterface)
492+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
493+
chat_target = MagicMock()
494+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
495+
496+
custom_prompt_path = _make_custom_system_prompt_yaml(tmp_path)
497+
custom_likert_path = _make_custom_scale_yaml(tmp_path, category="test_cat", min_val=1, max_val=5)
498+
499+
scorer = SelfAskLikertScorer(
500+
chat_target=chat_target,
501+
custom_likert_path=custom_likert_path,
502+
custom_system_prompt_path=custom_prompt_path,
503+
)
504+
505+
# The system prompt should come from the custom template, not the default one
506+
assert "Custom prompt for test_cat" in scorer._system_prompt
507+
assert "from 1 to 5" in scorer._system_prompt
508+
509+
510+
def test_custom_system_prompt_missing_params_rejected(tmp_path: Path):
511+
"""Verify that a custom system prompt missing required parameters raises ValueError."""
512+
memory = MagicMock(MemoryInterface)
513+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
514+
chat_target = MagicMock()
515+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
516+
517+
bad_prompt_path = _make_custom_system_prompt_yaml(tmp_path, include_all_params=False)
518+
custom_likert_path = _make_custom_scale_yaml(tmp_path)
519+
520+
with pytest.raises(ValueError, match="Custom system prompt YAML must define parameters"):
521+
SelfAskLikertScorer(
522+
chat_target=chat_target,
523+
custom_likert_path=custom_likert_path,
524+
custom_system_prompt_path=bad_prompt_path,
525+
)
526+
527+
528+
def test_both_likert_scale_and_custom_path_raises():
529+
"""Verify that providing both likert_scale and custom_likert_path raises ValueError."""
530+
memory = MagicMock(MemoryInterface)
531+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
532+
chat_target = MagicMock()
533+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
534+
535+
with pytest.raises(ValueError, match="Only one of"):
536+
SelfAskLikertScorer(
537+
chat_target=chat_target,
538+
likert_scale=LikertScalePaths.CYBER_SCALE,
539+
custom_likert_path=Path("dummy.yaml"),
540+
)
541+
542+
543+
def test_neither_likert_scale_nor_custom_path_raises():
544+
"""Verify that providing neither likert_scale nor custom_likert_path raises ValueError."""
545+
memory = MagicMock(MemoryInterface)
546+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
547+
chat_target = MagicMock()
548+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
549+
550+
with pytest.raises(ValueError, match="One of"):
551+
SelfAskLikertScorer(chat_target=chat_target)
552+
553+
554+
def test_custom_system_prompt_file_not_found():
555+
"""Verify that a non-existent custom_system_prompt_path raises FileNotFoundError."""
556+
memory = MagicMock(MemoryInterface)
557+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
558+
chat_target = MagicMock()
559+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
560+
561+
with pytest.raises(FileNotFoundError, match="Custom system prompt file not found"):
562+
SelfAskLikertScorer(
563+
chat_target=chat_target,
564+
likert_scale=LikertScalePaths.CYBER_SCALE,
565+
custom_system_prompt_path=Path("/does/not/exist.yaml"),
566+
)
567+
568+
569+
def test_custom_likert_yaml_not_a_dict_rejected(tmp_path: Path):
570+
"""Verify that a YAML file whose top-level structure is not a dict raises ValueError."""
571+
yaml_file = tmp_path / "bad_structure.yaml"
572+
yaml_file.write_text("- item1\n- item2\n", encoding="utf-8")
573+
574+
memory = MagicMock(MemoryInterface)
575+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
576+
chat_target = MagicMock()
577+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
578+
579+
with pytest.raises(ValueError, match="must contain a YAML mapping/dictionary"):
580+
SelfAskLikertScorer(chat_target=chat_target, custom_likert_path=yaml_file)
581+
582+
583+
def test_likert_scale_single_unique_value_rejected(tmp_path: Path):
584+
"""Verify that a scale with only one distinct score value raises ValueError."""
585+
yaml_file = tmp_path / "single_value.yaml"
586+
yaml_file.write_text(
587+
yaml.safe_dump(
588+
{
589+
"category": "test_harm",
590+
"scale_descriptions": [
591+
{"score_value": "3", "description": "Only level"},
592+
],
593+
}
594+
),
595+
encoding="utf-8",
596+
)
597+
598+
memory = MagicMock(MemoryInterface)
599+
with patch.object(CentralMemory, "get_memory_instance", return_value=memory):
600+
chat_target = MagicMock()
601+
chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget")
602+
603+
with pytest.raises(ValueError, match="at least two distinct score values"):
604+
SelfAskLikertScorer(chat_target=chat_target, custom_likert_path=yaml_file)

0 commit comments

Comments
 (0)