diff --git a/pyrit/score/response_handler.py b/pyrit/score/response_handler.py index 0c27efd120..7aa2c012c0 100644 --- a/pyrit/score/response_handler.py +++ b/pyrit/score/response_handler.py @@ -5,6 +5,7 @@ import abc import json +import math from abc import abstractmethod from collections.abc import Sequence from typing import TYPE_CHECKING @@ -146,7 +147,7 @@ class JsonSchemaResponseHandler(ResponseHandler): ``json.loads`` the text, then read the score value, rationale, optional description, category, and metadata from configurable keys. It also owns the response contract: the optional JSON schema handed to the target, and (when ``numeric_value`` is set) validating - that the parsed score value is numeric. + that the parsed score value is finite and numeric. """ def __init__( @@ -173,7 +174,8 @@ def __init__( should honor. Exposed via ``response_schema`` and forwarded to the target by the LLM round-trip. Defaults to None. numeric_value (bool): When True, ``parse`` requires the parsed score value to be - parsable as a float and raises ``InvalidJsonException`` otherwise. Defaults to False. + parsable as a finite float and raises ``InvalidJsonException`` otherwise. Defaults + to False. """ self._score_value_output_key = score_value_output_key self._rationale_output_key = rationale_output_key @@ -219,7 +221,7 @@ def parse( parsed category is not a string or a list of strings. InvalidJsonException: If the response is invalid JSON, is not a top-level JSON object, is missing a required key, or (when this handler is numeric) the score value is not - parsable as a float. + parsable as a finite float. """ response_json = remove_markdown_json(response_text) try: @@ -251,11 +253,15 @@ def parse( try: # A numeric handler requires the score value to be parsable as a float; a # well-formed-but-non-numeric value is treated as an invalid response. - float(score.raw_score_value) + parsed_value = float(score.raw_score_value) except ValueError: raise InvalidJsonException( message=f"Invalid JSON response, score_value should be a float not this: {score.raw_score_value}" ) from None + if not math.isfinite(parsed_value): + raise InvalidJsonException( + message=f"Invalid JSON response, score_value must be a finite float: {score.raw_score_value}" + ) return score diff --git a/tests/unit/score/test_response_handler.py b/tests/unit/score/test_response_handler.py index fac1b55ac0..df82970f8d 100644 --- a/tests/unit/score/test_response_handler.py +++ b/tests/unit/score/test_response_handler.py @@ -22,6 +22,18 @@ def test_json_schema_response_handler_rejects_non_object_response(response_text: ) +@pytest.mark.parametrize("score_value", ["nan", "NaN", "inf", "-inf", "Infinity"]) +def test_json_schema_response_handler_rejects_non_finite_numeric_values(score_value: str) -> None: + handler = JsonSchemaResponseHandler(numeric_value=True) + + with pytest.raises(InvalidJsonException, match="finite float"): + handler.parse( + response_text=f'{{"score_value": "{score_value}", "rationale": "test"}}', + scorer_identifier=SCORER_IDENTIFIER, + scored_prompt_id="test-id", + ) + + @pytest.mark.parametrize( ("json_value", "expected"), [