FIX Reject NaN in validate_temperature and validate_top_p

pyrit/prompt_target/common/utils.py

@@ -43,9 +43,9 @@ def validate_temperature(temperature: float | None) -> None:
         temperature: The temperature value to validate (0-2 inclusive).
 
     Raises:
-        PyritException: If temperature is not between 0 and 2 (inclusive).
+        PyritException: If temperature is NaN or not between 0 and 2 (inclusive).
     """
-    if temperature is not None and (temperature < 0 or temperature > 2):
+    if temperature is not None and not 0 <= temperature <= 2:
         raise PyritException(message="temperature must be between 0 and 2 (inclusive).")
 
 
@@ -57,9 +57,9 @@ def validate_top_p(top_p: float | None) -> None:
         top_p: The top_p value to validate (0-1 inclusive).
 
     Raises:
-        PyritException: If top_p is not between 0 and 1 (inclusive).
+        PyritException: If top_p is NaN or not between 0 and 1 (inclusive).
     """
-    if top_p is not None and (top_p < 0 or top_p > 1):
+    if top_p is not None and not 0 <= top_p <= 1:
         raise PyritException(message="top_p must be between 0 and 1 (inclusive).")
 
 

tests/unit/prompt_target/test_target_utils.py

@@ -64,6 +64,12 @@ def test_validate_temperature_above_two_raises():
         validate_temperature(2.1)
 
 
+@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf"), -float("inf")])
+def test_validate_temperature_non_finite_raises(invalid_value: float) -> None:
+    with pytest.raises(PyritException, match="temperature must be between 0 and 2"):
+        validate_temperature(invalid_value)
+
+
 def test_validate_top_p_none():
     validate_top_p(None)
 
@@ -90,6 +96,12 @@ def test_validate_top_p_above_one_raises():
         validate_top_p(1.1)
 
 
+@pytest.mark.parametrize("invalid_value", [float("nan"), float("inf"), -float("inf")])
+def test_validate_top_p_non_finite_raises(invalid_value: float) -> None:
+    with pytest.raises(PyritException, match="top_p must be between 0 and 1"):
+        validate_top_p(invalid_value)
+
+
 def test_validate_presence_penalty_none() -> None:
     validate_presence_penalty(None)