diff --git a/pyrit/score/float_scale/float_scale_score_aggregator.py b/pyrit/score/float_scale/float_scale_score_aggregator.py index f930dd69ae..e85cf026ab 100644 --- a/pyrit/score/float_scale/float_scale_score_aggregator.py +++ b/pyrit/score/float_scale/float_scale_score_aggregator.py @@ -60,12 +60,13 @@ def _create_aggregator( """ def aggregator(scores: Iterable[Score]) -> list[ScoreAggregatorResult]: - # Validate types and normalize input - for s in scores: + # Materialize before validating: `scores` is an Iterable, so validating by + # iterating it first would exhaust a generator and leave nothing to aggregate. + scores_list = list(scores) + for s in scores_list: if s.score_type != "float_scale": raise ValueError("All scores must be of type 'float_scale'.") - scores_list = list(scores) if not scores_list: if raise_on_empty: raise ValueError("No scores available for aggregation") @@ -182,12 +183,13 @@ def _create_aggregator_by_category( """ def aggregator(scores: Iterable[Score]) -> list[ScoreAggregatorResult]: - # Validate types and normalize input - for s in scores: + # Materialize before validating: `scores` is an Iterable, so validating by + # iterating it first would exhaust a generator and leave nothing to aggregate. + scores_list = list(scores) + for s in scores_list: if s.score_type != "float_scale": raise ValueError("All scores must be of type 'float_scale'.") - scores_list = list(scores) if not scores_list: # No scores; return a neutral result return [ diff --git a/pyrit/score/true_false/true_false_score_aggregator.py b/pyrit/score/true_false/true_false_score_aggregator.py index af97c9fdc8..a6527d7778 100644 --- a/pyrit/score/true_false/true_false_score_aggregator.py +++ b/pyrit/score/true_false/true_false_score_aggregator.py @@ -62,12 +62,13 @@ def _create_aggregator( """ def aggregator(scores: Iterable[Score]) -> ScoreAggregatorResult: - # Validate types and normalize input - for s in scores: + # Materialize before validating: `scores` is an Iterable, so validating by + # iterating it first would exhaust a generator and leave nothing to aggregate. + scores_list = list(scores) + for s in scores_list: if s.score_type != "true_false": raise ValueError("All scores must be of type 'true_false'.") - scores_list = list(scores) if not scores_list: # No scores; return a neutral result return ScoreAggregatorResult( diff --git a/tests/unit/score/test_float_scale_score_aggregator.py b/tests/unit/score/test_float_scale_score_aggregator.py index 19fac4bbff..4f26ee368d 100644 --- a/tests/unit/score/test_float_scale_score_aggregator.py +++ b/tests/unit/score/test_float_scale_score_aggregator.py @@ -2,6 +2,8 @@ # Licensed under the MIT license. +import pytest + from pyrit.models import ComponentIdentifier, Score from pyrit.score.float_scale.float_scale_score_aggregator import ( FloatScaleScoreAggregator, @@ -365,3 +367,48 @@ def test_average_raise_on_empty_with_no_scores(): with pytest.raises(ValueError, match="No scores available for aggregation"): FloatScaleScoreAggregator.AVERAGE_RAISE_ON_EMPTY([]) + + +def test_aggregators_accept_generators(): + """ + Aggregators are typed to take an Iterable, so a generator must aggregate the same + as the equivalent list. Validating by iterating before materializing exhausted the + generator and silently produced the empty-input result (0.0). + """ + values = [0.3, 0.9, 0.5] + + aggregators = [ + FloatScaleScoreAggregator.MAX, + FloatScaleScoreAggregator.MIN, + FloatScaleScoreAggregator.AVERAGE, + FloatScaleScorerByCategory.MAX, + FloatScaleScorerAllCategories.MAX, + ] + for aggregator in aggregators: + from_list = aggregator([_mk_score(v, category=["harm"]) for v in values]) + from_generator = aggregator(_mk_score(v, category=["harm"]) for v in values) + assert [r.value for r in from_generator] == [r.value for r in from_list] + + +def test_raise_on_empty_aggregator_accepts_generators(): + """A generator with scores must not trip the empty-input guard.""" + values = [0.3, 0.9, 0.5] + results = FloatScaleScoreAggregator.MAX_RAISE_ON_EMPTY(_mk_score(v) for v in values) + assert results[0].value == 0.9 + + +def test_generator_of_wrong_type_still_raises(): + """Materializing first must not weaken type validation.""" + bad = Score( + score_value="true", + score_value_description="", + score_type="true_false", + score_category=["test"], + score_rationale="", + score_metadata=None, + message_piece_id="1", + scorer_class_identifier=_TEST_SCORER_ID, + objective=None, + ) + with pytest.raises(ValueError, match="must be of type 'float_scale'"): + FloatScaleScoreAggregator.MAX(s for s in [bad]) diff --git a/tests/unit/score/test_true_false_score_aggregator.py b/tests/unit/score/test_true_false_score_aggregator.py index 8c64e4bfab..2ac4394c4f 100644 --- a/tests/unit/score/test_true_false_score_aggregator.py +++ b/tests/unit/score/test_true_false_score_aggregator.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import pytest + from pyrit.models import ComponentIdentifier, Score from pyrit.score import TrueFalseScoreAggregator @@ -236,3 +238,35 @@ def test_aggregator_single_score(): res = TrueFalseScoreAggregator.OR(scores) assert res.value is True assert res.rationale == "Single score rationale" + + +def test_aggregators_accept_generators(): + """ + Aggregators are typed to take an Iterable, so a generator must aggregate the same + as the equivalent list. Validating by iterating before materializing exhausted the + generator and silently produced the empty-input result (False). + """ + values = [False, True, False] + + for aggregator in (TrueFalseScoreAggregator.OR, TrueFalseScoreAggregator.AND): + from_list = aggregator([_mk_score(v, prr_id="1") for v in values]) + from_generator = aggregator(_mk_score(v, prr_id="1") for v in values) + assert from_generator.value == from_list.value + assert from_generator.description == from_list.description + + +def test_generator_of_wrong_type_still_raises(): + """Materializing first must not weaken type validation.""" + bad = Score( + score_value="0.5", + score_value_description="", + score_type="float_scale", + score_category=["test"], + score_rationale="", + score_metadata=None, + message_piece_id="1", + scorer_class_identifier=_TEST_SCORER_ID, + objective=None, + ) + with pytest.raises(ValueError, match="must be of type 'true_false'"): + TrueFalseScoreAggregator.OR(s for s in [bad])