From c16a4717175737456ede286eae48c748d510f552 Mon Sep 17 00:00:00 2001 From: Haseeb Nazir <36381672+iamhaseebn@users.noreply.github.com> Date: Sat, 22 Aug 2026 22:16:02 -0400 Subject: [PATCH] MAINT Decompose scorer execution policy and result finalization Signed-off-by: Haseeb Nazir <36381672+iamhaseebn@users.noreply.github.com> --- pyrit/score/message_scorer.py | 146 ++++++++++++++++------ tests/unit/score/test_message_scorer.py | 154 +++++++++++++++++++++++- 2 files changed, 264 insertions(+), 36 deletions(-) diff --git a/pyrit/score/message_scorer.py b/pyrit/score/message_scorer.py index 1a79272904..42bb8f7dbf 100644 --- a/pyrit/score/message_scorer.py +++ b/pyrit/score/message_scorer.py @@ -41,6 +41,20 @@ class MessageScoringOptions: skip_on_error_result: bool = False +@dataclass(frozen=True, kw_only=True) +class _PreparedMessageScoringInput: + """The effective input and policy decision for one message-scoring call.""" + + message: Message + expectation: ScoringExpectation | None + should_skip: bool = False + + @property + def objective(self) -> str | None: + """The objective attached to the effective expectation.""" + return self.expectation.objective if self.expectation else None + + def extract_objective_from_previous_turn(*, message: Message, memory: MemoryInterface) -> str: """ Read the text of the turn before an assistant message and use it as the objective. @@ -420,28 +434,53 @@ async def _score_resolved_message_async( PyritException: If scoring raises a PyRIT exception (re-raised with enhanced context). RuntimeError: If scoring raises a non-PyRIT exception (wrapped with scorer context). """ - objective = expectation.objective if expectation else None + scoring_input = self._prepare_message_scoring_input( + message=message, + expectation=expectation, + options=options, + infer_objective_from_request=infer_objective_from_request, + ) + if scoring_input.should_skip: + return [] - # Structured refusals are persisted as blocked error pieces, but scorers should - # receive the refusal explanation as text. Keep response_error="blocked" so - # refusal scorers can still use their deterministic blocked-response path. - scoring_message = self._apply_structured_refusal_substitution(message) + scores = await self._execute_message_scoring_async(scoring_input=scoring_input) + return self._finalize_message_scores(scoring_input=scoring_input, scores=scores) - # When score_blocked_content is enabled, blocked pieces with partial content - # take precedence and are replaced with text substitutes (response_error="none"). + def _prepare_message_scoring_input( + self, + *, + message: Message, + expectation: ScoringExpectation | None, + options: MessageScoringOptions, + infer_objective_from_request: bool, + ) -> _PreparedMessageScoringInput: + """ + Apply message policy and return the effective input for execution. + + Args: + message (Message): The acquired message. + expectation (ScoringExpectation | None): What to look for. + options (MessageScoringOptions): Message-only scoring policy. + infer_objective_from_request (bool): Whether to infer a missing objective. + + Returns: + _PreparedMessageScoringInput: The effective input and skip decision. + """ + objective = expectation.objective if expectation else None + scoring_message = self._apply_structured_refusal_substitution(message) if self.score_blocked_content: scoring_message = self._apply_blocked_content_substitution(scoring_message) self._validator.validate(scoring_message, objective=objective) + should_skip = False if options.role_filter is not None and message.get_piece().role != options.role_filter: logger.debug("Skipping scoring due to role filter mismatch.") - return [] - - if options.skip_on_error_result and self._should_skip_on_error(message): - return [] + should_skip = True + elif options.skip_on_error_result and self._should_skip_on_error(message): + should_skip = True - if infer_objective_from_request and (not objective): + if not should_skip and infer_objective_from_request and not objective: objective = extract_objective_from_previous_turn(message=message, memory=self._memory) effective_expectation = expectation @@ -453,44 +492,81 @@ async def _score_resolved_message_async( conditions=expectation.conditions, ) + return _PreparedMessageScoringInput( + message=scoring_message, + expectation=effective_expectation, + should_skip=should_skip, + ) + + async def _execute_message_scoring_async( + self, + *, + scoring_input: _PreparedMessageScoringInput, + ) -> list[Score]: + """ + Run scorer code with the scorer-layer exception policy. + + Args: + scoring_input (_PreparedMessageScoringInput): The prepared scoring input. + + Returns: + list[Score]: The scores produced by the scorer or its blocked fallback. + + Raises: + ScorerLLMResponseBlockedException: If the scorer's own LLM response is blocked + and ``raise_if_scorer_blocks`` is True. + PyritException: If scoring raises a PyRIT exception. + RuntimeError: If scoring raises a non-PyRIT exception. + """ try: - scores = await self._score_prepared_message_async( - message=scoring_message, - expectation=effective_expectation, + return await self._score_prepared_message_async( + message=scoring_input.message, + expectation=scoring_input.expectation, ) - except ScorerLLMResponseBlockedException as e: - # The scorer's own LLM response was content-filtered. By default this is a real - # error and propagates; when raise_if_scorer_blocks is False, fall back to the - # scorer's type default (False / 0.0) instead. The decision lives here in the - # scorer, not the transport (see doc/code/framework.md). + except ScorerLLMResponseBlockedException as error: if self.raise_if_scorer_blocks: - e.message = f"Error in scorer {self.__class__.__name__}: {e.message}" - e.args = (f"Status Code: {e.status_code}, Message: {e.message}",) + error.message = f"Error in scorer {self.__class__.__name__}: {error.message}" + error.args = (f"Status Code: {error.status_code}, Message: {error.message}",) raise logger.info( "Scorer %s LLM response was blocked by content filtering; " "returning default score (raise_if_scorer_blocks=False).", self.__class__.__name__, ) - scores = self._build_fallback_score( - message=scoring_message, - objective=objective, + return self._build_fallback_score( + message=scoring_input.message, + objective=scoring_input.objective, scorer_response_blocked=True, ) - except PyritException as e: - # Re-raise PyRIT exceptions with enhanced context while preserving type for retry decorators - e.message = f"Error in scorer {self.__class__.__name__}: {e.message}" - e.args = (f"Status Code: {e.status_code}, Message: {e.message}",) + except PyritException as error: + error.message = f"Error in scorer {self.__class__.__name__}: {error.message}" + error.args = (f"Status Code: {error.status_code}, Message: {error.message}",) raise - except Exception as e: - # Wrap non-PyRIT exceptions for better error tracing - raise RuntimeError(f"Error in scorer {self.__class__.__name__}: {str(e)}") from e + except Exception as error: + raise RuntimeError(f"Error in scorer {self.__class__.__name__}: {str(error)}") from error - if not scores and scoring_message.message_pieces: - scores = self._build_fallback_score(message=scoring_message, objective=objective) + def _finalize_message_scores( + self, + *, + scoring_input: _PreparedMessageScoringInput, + scores: list[Score], + ) -> list[Score]: + """ + Apply fallback behavior and clear links to ephemeral message pieces. - self._drop_ephemeral_score_links(message=scoring_message, scores=scores) + Args: + scoring_input (_PreparedMessageScoringInput): The prepared scoring input. + scores (list[Score]): The scores produced by execution. + Returns: + list[Score]: The finalized scores. + """ + if not scores and scoring_input.message.message_pieces: + scores = self._build_fallback_score( + message=scoring_input.message, + objective=scoring_input.objective, + ) + self._drop_ephemeral_score_links(message=scoring_input.message, scores=scores) return scores async def _score_prepared_message_async( diff --git a/tests/unit/score/test_message_scorer.py b/tests/unit/score/test_message_scorer.py index 99ca68ca95..bfda982a50 100644 --- a/tests/unit/score/test_message_scorer.py +++ b/tests/unit/score/test_message_scorer.py @@ -4,10 +4,11 @@ import dataclasses import inspect import uuid -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pyrit.exceptions import InvalidJsonException from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import ( ChatMessageRole, @@ -16,6 +17,7 @@ MatchesObjective, Message, MessagePiece, + PromptResponseError, Score, ScoringExpectation, ) @@ -306,6 +308,156 @@ async def test_no_expectation_means_no_objective(self): assert scorer.scored_objectives == [None] +@pytest.mark.usefixtures("patch_central_database") +class TestMessageScoringPolicyMatrix: + """The prepared input makes substitutions, filtering, and objective state explicit.""" + + @pytest.mark.parametrize( + ( + "role_filter", + "skip_on_error_result", + "structured_refusal", + "partial_content", + "score_blocked_content", + "input_response_error", + "expected_skip", + "expected_value", + "expected_response_error", + ), + [ + (None, False, None, None, False, "none", False, "response", "none"), + ("user", False, None, None, False, "none", True, "response", "none"), + (None, True, None, None, False, "blocked", True, "blocked", "blocked"), + (None, True, "I cannot assist.", None, False, "blocked", False, "I cannot assist.", "blocked"), + (None, True, None, "partial response", True, "blocked", False, "partial response", "none"), + (None, True, None, None, True, "blocked", True, "blocked", "blocked"), + ], + ) + def test_prepared_input_policy_matrix( + self, + role_filter: ChatMessageRole | None, + skip_on_error_result: bool, + structured_refusal: str | None, + partial_content: str | None, + score_blocked_content: bool, + input_response_error: PromptResponseError, + expected_skip: bool, + expected_value: str, + expected_response_error: PromptResponseError, + ) -> None: + is_blocked = input_response_error == "blocked" + piece = MessagePiece( + role="assistant", + original_value="blocked" if is_blocked else "response", + original_value_data_type="error" if is_blocked else "text", + converted_value_data_type="error" if is_blocked else "text", + response_error=input_response_error, + prompt_metadata={"partial_content": partial_content} if partial_content is not None else {}, + ) + if structured_refusal is not None: + piece.mark_as_structured_refusal(refusal=structured_refusal) + scorer = RecordingScorer() + scorer.score_blocked_content = score_blocked_content + + scoring_input = scorer._prepare_message_scoring_input( + message=piece.to_message(), + expectation=ScoringExpectation(objective="objective"), + options=MessageScoringOptions( + role_filter=role_filter, + skip_on_error_result=skip_on_error_result, + ), + infer_objective_from_request=False, + ) + + assert scoring_input.should_skip is expected_skip + assert scoring_input.objective == "objective" + assert scoring_input.message.get_value() == expected_value + assert scoring_input.message.get_piece().response_error == expected_response_error + + def test_inferred_objective_is_part_of_the_prepared_input(self, sqlite_instance: MemoryInterface) -> None: + conversation_id = str(uuid.uuid4()) + sqlite_instance.add_message_to_memory( + request=MessagePiece( + role="user", + original_value="the inferred objective", + conversation_id=conversation_id, + sequence=0, + ).to_message() + ) + message = _assistant_message("response", conversation_id=conversation_id) + scorer = RecordingScorer() + + scoring_input = scorer._prepare_message_scoring_input( + message=message, + expectation=None, + options=MessageScoringOptions(), + infer_objective_from_request=True, + ) + + assert scoring_input.objective == "the inferred objective" + assert scoring_input.expectation == ScoringExpectation(objective="the inferred objective") + + +@pytest.mark.usefixtures("patch_central_database") +class TestMessageScoringExecutionAndFinalization: + """Execution policy and outer score persistence retain their exact contracts.""" + + async def test_pyrit_exception_keeps_type_and_gains_scorer_context(self) -> None: + scorer = RecordingScorer() + error = InvalidJsonException(message="invalid scorer response") + + with patch.object( + scorer, + "_score_prepared_message_async", + new=AsyncMock(side_effect=error), + ): + with pytest.raises(InvalidJsonException) as exception_info: + await scorer.score_async(scorable=MessageScorable.from_message(_assistant_message())) + + assert exception_info.value is error + assert error.message == "Error in scorer RecordingScorer: invalid scorer response" + + async def test_non_pyrit_exception_is_wrapped_with_scorer_context(self) -> None: + scorer = RecordingScorer() + error = ValueError("unexpected scorer failure") + + with patch.object( + scorer, + "_score_prepared_message_async", + new=AsyncMock(side_effect=error), + ): + with pytest.raises(RuntimeError, match="Error in scorer RecordingScorer") as exception_info: + await scorer.score_async(scorable=MessageScorable.from_message(_assistant_message())) + + assert exception_info.value.__cause__ is error + + async def test_success_validates_and_persists_scores_once(self) -> None: + scorer = RecordingScorer() + with ( + patch.object(scorer, "validate_return_scores", wraps=scorer.validate_return_scores) as validate, + patch.object(scorer._memory, "add_scores_to_memory") as persist, + ): + scores = await scorer.score_async(scorable=MessageScorable.from_message(_assistant_message())) + + validate.assert_called_once_with(scores=scores) + persist.assert_called_once_with(scores=scores) + + async def test_skipped_input_is_not_validated_or_persisted(self) -> None: + scorer = RecordingScorer() + with ( + patch.object(scorer, "validate_return_scores", wraps=scorer.validate_return_scores) as validate, + patch.object(scorer._memory, "add_scores_to_memory") as persist, + ): + scores = await scorer.score_async( + scorable=MessageScorable.from_message(_assistant_message()), + message_options=MessageScoringOptions(role_filter="user"), + ) + + assert scores == [] + validate.assert_not_called() + persist.assert_not_called() + + @pytest.mark.usefixtures("patch_central_database") class TestDeprecatedParameters: """The legacy message-shaped parameters survive one release behind a warning."""