diff --git a/pyrit/memory/analytics_sql.py b/pyrit/memory/analytics_sql.py index 84a8d65519..6779641b09 100644 --- a/pyrit/memory/analytics_sql.py +++ b/pyrit/memory/analytics_sql.py @@ -49,6 +49,24 @@ class UnicodeLower(FunctionElement[str]): inherit_cache = True +def _sqlite_fold_case(*, value: str, fallback_function: str) -> str: + """ + Keep ASCII text in SQLite as plain ``lower`` and use a registered fold otherwise. + + Byte/character lengths differ for Unicode and embedded NULs. Non-text + values keep the Python conversion rather than SQLite's stringification. + Both Unicode fold elements share this expression so the BLOB-length guard + and the ASCII fast path cannot drift apart. + + Returns: + str: The compiled fold expression. + """ + return ( + f"CASE WHEN typeof({value}) = 'text' AND length(CAST({value} AS BLOB)) = length({value}) " + f"THEN lower({value}) ELSE {fallback_function}({value}) END" + ) + + @compiles(UnicodeLower, "sqlite") def _sqlite_unicode_lower(element: UnicodeLower, compiler: Any, **kwargs: Any) -> str: """ @@ -58,12 +76,47 @@ def _sqlite_unicode_lower(element: UnicodeLower, compiler: Any, **kwargs: Any) - str: Exact Unicode-aware lowercasing without a Python/GIL callback for ASCII text. """ value = _arguments(element, compiler, **kwargs)[0] - # Byte/character lengths differ for Unicode and embedded NULs. Non-text - # values keep the existing Python conversion rather than SQLite's stringification. - return ( - f"CASE WHEN typeof({value}) = 'text' AND length(CAST({value} AS BLOB)) = length({value}) " - f"THEN lower({value}) ELSE {UnicodeLower.SQLITE_FUNCTION_NAME}({value}) END" - ) + return _sqlite_fold_case(value=value, fallback_function=UnicodeLower.SQLITE_FUNCTION_NAME) + + +class UnicodeCaseFold(FunctionElement[str]): + """ + Fold text per character with Python casefold semantics for exact matching. + + Unlike ``UnicodeLower``, every character maps on its own, so verbatim + text always matches regardless of context sensitive forms such as Greek + final sigma. Full casefold semantics, including ``ss`` matching the sharp + s, hold on SQLite through the registered function; SQL Server compiles to + plain ``LOWER``, where such expansions follow the column collation. + """ + + SQLITE_FUNCTION_NAME: ClassVar[str] = "pyrit_analytics_casefold" + type = UnicodeText() + inherit_cache = True + + +@compiles(UnicodeCaseFold, "sqlite") +def _sqlite_unicode_casefold(element: UnicodeCaseFold, compiler: Any, **kwargs: Any) -> str: + """ + Reuse the shared SQLite fold expression with the registered casefold. + + Returns: + str: The compiled per character fold expression. + """ + value = _arguments(element, compiler, **kwargs)[0] + return _sqlite_fold_case(value=value, fallback_function=UnicodeCaseFold.SQLITE_FUNCTION_NAME) + + +@compiles(UnicodeCaseFold, "mssql") +def _mssql_unicode_casefold(element: UnicodeCaseFold, compiler: Any, **kwargs: Any) -> str: + """ + Compile the casefold for SQL Server, whose LOWER stays per character. + + Returns: + str: The compiled per character fold expression. + """ + value = _arguments(element, compiler, **kwargs)[0] + return f"LOWER({value})" @compiles(UnicodeLower, "mssql") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index e230165ae0..937557bdf0 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -27,6 +27,7 @@ Select, String, Unicode, + UnicodeText, and_, case, column, @@ -58,7 +59,15 @@ from pyrit.common.async_compatibility import legacy_sync_override, run_legacy_sync_async from pyrit.common.deprecation import print_deprecation_message from pyrit.common.pagination import DecodedKeysetCursor -from pyrit.memory.analytics_sql import JsonArrayItems, JsonArrayJoin, JsonArrayValue, JsonContainer, JsonScalar, UuidHex +from pyrit.memory.analytics_sql import ( + JsonArrayItems, + JsonArrayJoin, + JsonArrayValue, + JsonContainer, + JsonScalar, + UnicodeCaseFold, + UuidHex, +) if TYPE_CHECKING: from pyrit.memory.memory_embedding import MemoryEmbedding @@ -1062,7 +1071,12 @@ async def get_findings_async( if title_query: # SQL Server also treats "[" as a LIKE wildcard. escaped_title = title_query.replace("/", "//").replace("%", "/%").replace("_", "/_").replace("[", "/[") - statement = statement.where(FindingEntry.title.ilike(f"%{escaped_title}%", escape="/")) + # SQLite's lower() only folds ASCII, so both sides fold through the + # per character UnicodeCaseFold, matching the seed value search. + pattern = "%" + escaped_title + "%" + statement = statement.where( + UnicodeCaseFold(FindingEntry.title).like(UnicodeCaseFold(literal(pattern, UnicodeText())), escape="/") + ) statement = statement.order_by(*self._finding_order()).limit(limit).offset(offset) async with await self.get_session_async() as session: return [entry.get_finding() for entry in (await session.scalars(statement)).all()] @@ -5477,12 +5491,17 @@ def _seed_example_filters( member_conditions.append(SeedEntry.seed_type.in_(list(seed_types))) if value_search: # A simulated-conversation value is JSON, so its keys would match common words such as "prompt". + # casefold maps each character independently, so lowering the term separately from the + # stored text cannot lose context sensitive forms such as Greek final sigma, and verbatim + # text always still matches. pattern = "%" + re.sub(r"([\\%_\[])", r"\\\1", value_search) + "%" member_conditions.append( and_( SeedEntry.data_type == "text", SeedEntry.seed_type != "simulated_conversation", - SeedEntry.value.ilike(pattern, escape="\\"), + UnicodeCaseFold(SeedEntry.value).like( + UnicodeCaseFold(literal(pattern, UnicodeText())), escape="\\" + ), ) ) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 196ce51808..f11c3df499 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -28,7 +28,7 @@ from pyrit.common.path import DB_DATA_PATH from pyrit.common.singleton import Singleton -from pyrit.memory.analytics_sql import UnicodeLower +from pyrit.memory.analytics_sql import UnicodeCaseFold, UnicodeLower from pyrit.memory.memory_interface import MemoryInterface from pyrit.memory.memory_models import ( AttackResultEntry, @@ -280,15 +280,28 @@ def _unicode_lower(value: str | int | float | bytes | None) -> str | None: """ return str(value).lower() if value is not None else None + @staticmethod + def _unicode_casefold(value: str | int | float | bytes | None) -> str | None: + """ + Casefold SQLite text per character, retaining NULL semantics. + + Returns: + str | None: The folded value, or SQL NULL. + """ + return str(value).casefold() if value is not None else None + @staticmethod def _register_analytics_lower(*, engine: Engine) -> None: - """Install the Unicode lowercase function for analytics and JSON filters on every pooled SQLite connection.""" + """Install the Unicode fold functions for analytics and JSON filters on every pooled SQLite connection.""" @event.listens_for(engine, "connect") def register(dbapi_connection: Any, connection_record: Any) -> None: dbapi_connection.create_function( UnicodeLower.SQLITE_FUNCTION_NAME, 1, SQLiteMemory._unicode_lower, deterministic=True ) + dbapi_connection.create_function( + UnicodeCaseFold.SQLITE_FUNCTION_NAME, 1, SQLiteMemory._unicode_casefold, deterministic=True + ) async def get_session_async(self) -> AsyncSession: """ diff --git a/tests/unit/memory/memory_interface/test_interface_seed_examples.py b/tests/unit/memory/memory_interface/test_interface_seed_examples.py index f7c4700c26..4985aad54c 100644 --- a/tests/unit/memory/memory_interface/test_interface_seed_examples.py +++ b/tests/unit/memory/memory_interface/test_interface_seed_examples.py @@ -161,6 +161,48 @@ async def test_get_seed_examples_harm_filter_folds_non_ascii_case(sqlite_instanc assert await _ids(sqlite_instance, harm_categories=["Überwachung"]) == [surveillance] +async def test_get_seed_examples_value_search_folds_non_ascii_case(sqlite_instance: MemoryInterface): + cafe, wildcard = uuid4(), uuid4() + await _add( + sqlite_instance, + SeedPrompt(value="CAFÉ au lait tactics", dataset_name=DATASET, prompt_group_id=cafe), + SeedPrompt(value="50%_off coupon", dataset_name=DATASET, prompt_group_id=wildcard), + ) + + assert await _ids(sqlite_instance, value_search="CAFÉ") == [cafe] + assert await _ids(sqlite_instance, value_search="café") == [cafe] + assert await _ids(sqlite_instance, value_search="Au Lait") == [cafe] + assert await _ids(sqlite_instance, value_search="50%_off") == [wildcard] + + +async def test_get_seed_examples_value_search_matches_greek_final_sigma(sqlite_instance: MemoryInterface): + """Verbatim Greek text must match regardless of context sensitive final sigma.""" + astro = uuid4() + await _add( + sqlite_instance, + SeedPrompt(value="ΑΣΤΡΟ navigation charts", dataset_name=DATASET, prompt_group_id=astro), + ) + + # 'ΑΣ' lower()s to 'ας' while the stored text lower()s to 'αστρο', so naive + # lowering of the two sides separately loses the match; casefold keeps it. + assert await _ids(sqlite_instance, value_search="ΑΣ") == [astro] + assert await _ids(sqlite_instance, value_search="ασ") == [astro] + + +async def test_get_seed_examples_value_search_casefold_expansions_hold_on_sqlite( + sqlite_instance: MemoryInterface, +): + """SQLite's registered casefold maps the sharp s to ss, so both spellings match.""" + strasse = uuid4() + await _add( + sqlite_instance, + SeedPrompt(value="Straße crossing signs", dataset_name=DATASET, prompt_group_id=strasse), + ) + + assert await _ids(sqlite_instance, value_search="STRASSE") == [strasse] + assert await _ids(sqlite_instance, value_search="straße") == [strasse] + + async def test_get_seed_examples_keeps_named_and_unnamed_scopes_apart(sqlite_instance: MemoryInterface): group = uuid4() named = SeedPrompt(value="named", dataset_name=DATASET, prompt_group_id=group) diff --git a/tests/unit/memory/test_analytics_migration.py b/tests/unit/memory/test_analytics_migration.py index 5bbe0c531b..b6251f625e 100644 --- a/tests/unit/memory/test_analytics_migration.py +++ b/tests/unit/memory/test_analytics_migration.py @@ -15,13 +15,14 @@ from alembic.config import Config from alembic.operations import Operations from alembic.script import ScriptDirectory -from sqlalchemy import String, create_engine, inspect, text +from sqlalchemy import String, UnicodeText, create_engine, inspect, literal, text from sqlalchemy.dialects import mssql, sqlite from sqlalchemy.engine import Connection from sqlalchemy.schema import CreateIndex, CreateTable from pyrit.memory.alembic.versions.e5f7a9c1b3d2_add_identifiers_tables import IdentifierGraphInserter from pyrit.memory.analytics_identity_v1 import ObjectiveTargetAnalyticsIdentityV1 +from pyrit.memory.analytics_sql import UnicodeCaseFold from pyrit.memory.memory_models import AttackResultEntry from pyrit.models import AtomicAttackIdentifier, AttackIdentifier, TargetIdentifier, compute_eval_hash @@ -324,6 +325,43 @@ def test_analytics_migration_preserves_outcome_at_the_index_key_limit(tmp_path: engine.dispose() +def test_unicode_casefold_compiles_per_dialect() -> None: + """SQLite keeps the registered casefold fallback and SQL Server uses LOWER exactly.""" + expression = UnicodeCaseFold(literal("CAFÉ", UnicodeText())) + sqlite_sql = str(expression.compile(dialect=sqlite.dialect())) + mssql_sql = str(expression.compile(dialect=mssql.dialect())) + assert "pyrit_analytics_casefold" in sqlite_sql + assert "lower(" in sqlite_sql # the ASCII fast path stays plain lower + # Pinned exactly: a wrapper around the fallback would change Azure folding. + assert mssql_sql == "LOWER(:param_1)" + + +def test_seed_value_search_folds_both_sides_in_the_real_statement() -> None: + """The seed search statement itself folds column and pattern in SQL.""" + from sqlalchemy import true + + from pyrit.memory.sqlite_memory import SQLiteMemory + + conditions = SQLiteMemory._seed_example_filters( + SQLiteMemory.__new__(SQLiteMemory), + scope=true(), + data_types=None, + harm_categories=None, + seed_types=None, + value_search="CAFÉ", + ) + assert conditions, "the value search must produce a condition" + for dialect, fold_marker in ( + (sqlite.dialect(), "pyrit_analytics_casefold"), + (mssql.dialect(), "LOWER"), + ): + for condition in conditions: + compiled = str(condition.compile(dialect=dialect)) + # Column and pattern both fold in SQL, so neither backend can mix + # foldings by lowering one side in Python. + assert compiled.count(fold_marker) == 2, compiled + + @pytest.mark.parametrize("dialect", [sqlite.dialect(), mssql.dialect()]) def test_generated_lookup_and_index_ddl_are_dialect_specific(dialect: Dialect) -> None: sql = str(CreateTable(AttackResultEntry.__table__).compile(dialect=dialect)) diff --git a/tests/unit/memory/test_findings.py b/tests/unit/memory/test_findings.py index 94d45c0e7a..de14efa854 100644 --- a/tests/unit/memory/test_findings.py +++ b/tests/unit/memory/test_findings.py @@ -64,6 +64,24 @@ async def test_update_operation_keeps_names_unique_and_valid(sqlite_instance: SQ assert await sqlite_instance.get_operations_async() == [other, saved] +async def test_get_findings_title_search_folds_non_ascii_case(sqlite_instance: SQLiteMemory) -> None: + """The title search matches non-ASCII titles in both fold directions and keeps wildcards literal.""" + from pyrit.models import Finding, Operation + + operation = await sqlite_instance.add_operation_async(Operation(name="Search")) + await sqlite_instance.add_finding_async( + finding=Finding(operation_id=operation.id, title="CAF\u00c9 security review", severity="informational") + ) + await sqlite_instance.add_finding_async( + finding=Finding(operation_id=operation.id, title="100% reviewed", severity="low") + ) + + by_fold = await sqlite_instance.get_findings_async(operation_id=operation.id, title_query="caf\u00e9") + assert [finding.title for finding in by_fold] == ["CAF\u00c9 security review"] + literal = await sqlite_instance.get_findings_async(operation_id=operation.id, title_query="100%") + assert [finding.title for finding in literal] == ["100% reviewed"] + + async def test_add_finding_rejects_a_missing_operation(sqlite_instance: SQLiteMemory) -> None: with pytest.raises(LookupError, match="Operation"): await sqlite_instance.add_finding_async(finding=Finding(operation_id=uuid4(), title="Orphan", severity="low")) @@ -538,7 +556,10 @@ async def test_finding_search_escapes_sql_server_bracket_wildcards(sqlite_instan await sqlite_instance.get_findings_async(title_query="[Proof]") statement = session.scalars.call_args.args[0] compiled = statement.compile(dialect=mssql.dialect()) - assert compiled.params["title_1"] == "%/[Proof]%" + # The pattern folds in SQL on both sides, so the bound parameter carries + # the raw escaped title and the fold wrappers appear in the SQL itself. + assert "%/[Proof]%" in compiled.params.values() + assert "LOWER(" in str(compiled) assert "ESCAPE '/'" in str(compiled)