From dc3a1a289192b10dae462a656c6ed4db63415dc2 Mon Sep 17 00:00:00 2001 From: JR Date: Wed, 28 Oct 2026 23:00:13 +0100 Subject: [PATCH] FIX: match one category in get_scores score_category filter --- pyrit/memory/memory_interface.py | 18 +++++--- .../memory_interface/test_interface_scores.py | 44 +++++++++++++++++++ 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index e230165ae0..89e878ef47 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -3502,7 +3502,7 @@ def _execute_get_scores( *, score_ids: Sequence[str] | None = None, score_type: str | None = None, - score_category: str | None = None, + score_category: str | Sequence[str] | None = None, sent_after: datetime | None = None, sent_before: datetime | None = None, identifier_filters: Sequence[IdentifierFilter] | None = None, @@ -3514,7 +3514,8 @@ def _execute_get_scores( Args: score_ids (Sequence[str] | None): A list of score IDs to filter by. score_type (str | None): The type of the score to filter by. - score_category (str | None): The category of the score to filter by. + score_category (str | Sequence[str] | None): Return scores that list this category, or every category + of a sequence. Names match whole and case-insensitively. sent_after (datetime | None): Filter for scores sent after this datetime. sent_before (datetime | None): Filter for scores sent before this datetime. identifier_filters (Sequence[IdentifierFilter] | None): A sequence of IdentifierFilter objects that @@ -3533,7 +3534,11 @@ def _execute_get_scores( if score_type: conditions.append(ScoreEntry.score_type == score_type) if score_category: - conditions.append(ScoreEntry.score_category == score_category) + # score_category is stored as a JSON list, so match one whole element of it. + categories = [score_category] if isinstance(score_category, str) else list(score_category) + self._add_list_conditions( + field=ScoreEntry.score_category, conditions=conditions, values=categories, exact=True + ) if sent_after: conditions.append(ScoreEntry.timestamp >= sent_after) if sent_before: @@ -9692,7 +9697,7 @@ def get_scores( *, score_ids: Sequence[str] | None = None, score_type: str | None = None, - score_category: str | None = None, + score_category: str | Sequence[str] | None = None, sent_after: datetime | None = None, sent_before: datetime | None = None, identifier_filters: Sequence[IdentifierFilter] | None = None, @@ -9725,7 +9730,7 @@ async def get_scores_async( *, score_ids: Sequence[str] | None = None, score_type: str | None = None, - score_category: str | None = None, + score_category: str | Sequence[str] | None = None, sent_after: datetime | None = None, sent_before: datetime | None = None, identifier_filters: Sequence[IdentifierFilter] | None = None, @@ -9737,7 +9742,8 @@ async def get_scores_async( Args: score_ids (Sequence[str] | None): A list of score IDs to filter by. score_type (str | None): The type of the score to filter by. - score_category (str | None): The category of the score to filter by. + score_category (str | Sequence[str] | None): Return scores that list this category, or every category + of a sequence. Names match whole and case-insensitively. sent_after (datetime | None): Filter for scores sent after this datetime. sent_before (datetime | None): Filter for scores sent before this datetime. identifier_filters (Sequence[IdentifierFilter] | None): A sequence of IdentifierFilter objects that diff --git a/tests/unit/memory/memory_interface/test_interface_scores.py b/tests/unit/memory/memory_interface/test_interface_scores.py index 98fa71e074..839d2e80d7 100644 --- a/tests/unit/memory/memory_interface/test_interface_scores.py +++ b/tests/unit/memory/memory_interface/test_interface_scores.py @@ -448,3 +448,47 @@ async def test_get_scores_by_scorer_identifier_filter( ], ) assert len(results) == 0 + + +async def test_get_scores_by_score_category( + sqlite_instance: MemoryInterface, + sample_conversation_entries: Sequence[PromptMemoryEntry], +): + prompt_id = sample_conversation_entries[0].id + sqlite_instance._insert_entries(entries=sample_conversation_entries) + + def make_score(value: str, categories: list[str] | None, score_type: str = "float_scale") -> Score: + return Score( + score_value=value, + score_value_description="", + score_type=score_type, + score_category=categories, + score_rationale="", + score_metadata={}, + scorer_class_identifier=_test_scorer_id(), + message_piece_id=prompt_id, + ) + + await sqlite_instance.add_scores_to_memory_async( + scores=[ + make_score("0.1", ["hate"]), + make_score("0.2", ["violence", "hate"]), + make_score("0.3", ["hate_speech"]), + make_score("0.4", ["HaTe"]), + make_score("0.5", None), + make_score("0.6", []), + make_score("true", ["hate"], score_type="true_false"), + ] + ) + + async def values_for(category: str | list[str], score_type: str | None = None): + scores = await sqlite_instance.get_scores_async(score_category=category, score_type=score_type) + return sorted(score.score_value for score in scores) + + assert await values_for("hate") == ["0.1", "0.2", "0.4", "true"] + assert await values_for("HATE", score_type="float_scale") == ["0.1", "0.2", "0.4"] + assert await values_for("violence") == ["0.2"] + assert await values_for("hate_speech") == ["0.3"] + assert await values_for("speech") == [] + assert await values_for(["hate", "violence"]) == ["0.2"] + assert await values_for(["hate"], score_type="float_scale") == ["0.1", "0.2", "0.4"]