Merge branch 'main' into benchmark

Author: ValbuenaVC

.github/workflows/build_and_test.yml

@@ -31,14 +31,14 @@ jobs:
     permissions:
       contents: read
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v5
 
-      - uses: actions/setup-python@v3
+      - uses: actions/setup-python@v6
         with:
           python-version: ${{ env.PRE_COMMIT_PYTHON_VERSION }}
 
       - name: Cache pre-commit environments
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: ~/.cache/pre-commit
           key: pre-commit-${{ runner.os }}-${{ hashFiles('.pre-commit-config.yaml') }}
@@ -83,7 +83,7 @@ jobs:
       fail-fast: true
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
-        python: ["3.10", "3.11", "3.12", "3.13"]
+        python: ["3.10", "3.11", "3.12", "3.13", "3.14"]
         package_name: ["pyrit"]
         package_extras: ["dev", "dev_all"]
     runs-on: ${{ matrix.os }}
@@ -94,10 +94,10 @@ jobs:
       checks: write
       pull-requests: write
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v5
 
       # Set up Python
-      - uses: actions/setup-python@v3
+      - uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python }}
 

.pre-commit-config.yaml

@@ -30,6 +30,18 @@ repos:
         files: ^(doc/.*\.(py|ipynb|md)|doc/myst\.yml)$
         pass_filenames: false
         additional_dependencies: ['pyyaml']
+      - id: enforce_alembic_revision_immutability
+        name: Enforce Alembic Revision Immutability
+        entry: python ./build_scripts/enforce_alembic_revision_immutability.py
+        language: python
+        files: ^pyrit/memory/alembic/versions/.*\.py$
+        pass_filenames: false
+      - id: memory-migrations-check
+        name: Check Memory Migrations
+        entry: python ./build_scripts/memory_migrations.py check
+        language: system
+        pass_filenames: false
+        files: ^pyrit/memory/(memory_models\.py|alembic/.*|migration\.py)$
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0

build_scripts/enforce_alembic_revision_immutability.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Migration history must be immutable. This hook enforces that by preventing deletion or updates to migration scripts.
+
+Checks both staged changes (local pre-commit) and the full branch diff against origin/main (CI).
+"""
+
+import subprocess
+import sys
+
+_VERSIONS_PATH = "pyrit/memory/alembic/versions/"
+
+
+def _git(*args: str) -> str:
+    result = subprocess.run(["git", *args], capture_output=True, text=True)
+    return result.stdout.strip()
+
+
+def _has_non_add_changes(diff_spec: list[str]) -> bool:
+    output = _git("diff", "--name-status", *diff_spec, "--", _VERSIONS_PATH)
+    return any(line and not line.startswith("A") for line in output.splitlines())
+
+
+def has_revision_violations() -> bool:
+    # Local pre-commit: check staged changes
+    if _has_non_add_changes(["--cached"]):
+        return True
+
+    # CI: check full branch diff against origin/main
+    merge_base = _git("merge-base", "origin/main", "HEAD")
+    return bool(merge_base and _has_non_add_changes([f"{merge_base}...HEAD"]))
+
+
+if __name__ == "__main__":
+    if has_revision_violations():
+        print("[ERROR] Migration scripts can only be added, not modified or deleted.")
+        sys.exit(1)

build_scripts/memory_migrations.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import argparse
+import sys
+import tempfile
+from pathlib import Path
+
+from alembic.util.exc import AutogenerateDiffsDetected
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine
+
+from pyrit.memory.migration import check_schema_migrations, generate_schema_migration, run_schema_migrations
+
+# ANSI color codes
+_RED = "\033[91m"
+_RESET = "\033[0m"
+
+
+def _print_error(message: str) -> None:
+    """Print an error message in red to stderr."""
+    print(f"{_RED}{message}{_RESET}", file=sys.stderr)
+
+
+def _create_temp_engine() -> tuple[Engine, Path]:
+    """Create a temp SQLite database upgraded to head and return engine and path."""
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
+        tmp_path = Path(tmp.name)
+    engine = create_engine(f"sqlite:///{tmp_path}")
+    run_schema_migrations(engine=engine)
+    return engine, tmp_path
+
+
+def _cmd_generate(*, message: str, force: bool = False) -> None:
+    """Generate a new Alembic revision from model changes."""
+    engine, tmp_path = _create_temp_engine()
+    try:
+        generate_schema_migration(engine=engine, message=message, force=force)
+        print("Migration file generated. Review it carefully before committing.")
+    except RuntimeError as e:
+        _print_error(str(e))
+        raise SystemExit(1) from e
+    finally:
+        engine.dispose()
+        tmp_path.unlink(missing_ok=True)
+
+
+def _cmd_check() -> None:
+    """Verify all migrations apply cleanly and schema matches models."""
+    engine, tmp_path = _create_temp_engine()
+    try:
+        check_schema_migrations(engine=engine)
+    except AutogenerateDiffsDetected as e:
+        _print_error(f"Migration check failed. Run 'generate' to create a migration. Error: {e}")
+        raise SystemExit(1) from e
+    finally:
+        engine.dispose()
+        tmp_path.unlink(missing_ok=True)
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    """Build the CLI argument parser."""
+    parser = argparse.ArgumentParser(
+        description="PyRIT memory migration tool. Generate and validate migrations based on the current memory models."
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    gen = sub.add_parser("generate", help="Generate a new migration from model changes.")
+    gen.add_argument("-m", "--message", required=True, help="Migration message.")
+    gen.add_argument("--force", action="store_true", help="Generate migration even if no changes detected.")
+
+    sub.add_parser("check", help="Verify all migrations apply cleanly and add up to the current memory models.")
+
+    return parser
+
+
+def main() -> int:
+    """Dispatch the selected migration command."""
+    args = _build_parser().parse_args()
+
+    if args.command == "generate":
+        _cmd_generate(message=args.message, force=args.force)
+    elif args.command == "check":
+        _cmd_check()
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

doc/code/datasets/1_loading_datasets.ipynb

@@ -38,6 +38,7 @@
     "TDC23 [@mazeika2023tdc],\n",
     "ToxicChat [@lin2023toxicchat],\n",
     "VLSU [@palaskar2025vlsu],\n",
+    "VLGuard [@zong2024vlguard],\n",
     "XSTest [@rottger2023xstest],\n",
     "AILuminate [@vidgen2024ailuminate],\n",
     "Transphobia Awareness [@scheuerman2025transphobia],\n",

doc/code/datasets/1_loading_datasets.py

@@ -42,6 +42,7 @@
 # TDC23 [@mazeika2023tdc],
 # ToxicChat [@lin2023toxicchat],
 # VLSU [@palaskar2025vlsu],
+# VLGuard [@zong2024vlguard],
 # XSTest [@rottger2023xstest],
 # AILuminate [@vidgen2024ailuminate],
 # Transphobia Awareness [@scheuerman2025transphobia],

doc/contributing/11_memory_models.md

@@ -0,0 +1,103 @@
+# Memory Models & Migrations
+
+This guide covers how to work with PyRIT's memory models — where they live, how to add or update them, and how the migration system works.
+
+## Where Things Live
+
+| What | Path |
+|---|---|
+| ORM models (SQLAlchemy) | `pyrit/memory/memory_models.py` |
+| Domain objects they map to | `pyrit/models/` (e.g. `MessagePiece`, `Score`, `Seed`, `AttackResult`, `ScenarioResult`) |
+| Alembic migration environment | `pyrit/memory/alembic/env.py` |
+| Migration revisions | `pyrit/memory/alembic/versions/` |
+| Migration helpers | `pyrit/memory/migration.py` |
+| CLI migration tool | `build_scripts/memory_migrations.py` |
+| Schema diagram | `doc/code/memory/10_schema_diagram.md` |
+
+## Current Models
+
+All models inherit from the SQLAlchemy `Base` declarative class and live in `memory_models.py`:
+
+- **`PromptMemoryEntry`** — prompt/response data (`PromptMemoryEntries` table)
+- **`ScoreEntry`** — evaluation results (`ScoreEntries` table)
+- **`EmbeddingDataEntry`** — embeddings for semantic search (`EmbeddingData` table)
+- **`SeedEntry`** — dataset prompts/templates (`SeedPromptEntries` table)
+- **`AttackResultEntry`** — attack execution results (`AttackResultEntries` table)
+- **`ScenarioResultEntry`** — scenario execution metadata (`ScenarioResultEntries` table)
+
+Each entry model has a corresponding domain object and conversion methods (e.g. `PromptMemoryEntry.__init__(entry: MessagePiece)` and `get_message_piece()`).
+
+## Adding or Updating a Model
+
+### 1. Edit the model
+
+Make your changes in `pyrit/memory/memory_models.py`. Follow these conventions:
+
+- Use `mapped_column()` with explicit types.
+- Use `CustomUUID` for all UUID columns (handles cross-database compatibility).
+- Add foreign keys where relationships exist.
+- Include `pyrit_version` on new entry models.
+
+### 2. Generate a migration
+
+```bash
+python build_scripts/memory_migrations.py generate -m "short description of change"
+```
+
+This creates a new revision file under `pyrit/memory/alembic/versions/`. **Review the generated file carefully** — auto-generated migrations may need manual adjustments (e.g. for data migrations or default values).
+
+### 3. Validate the migration
+
+```bash
+python build_scripts/memory_migrations.py check
+```
+
+This verifies the schema produced by running all migrations matches the current models. Both pre-commit hooks (see below) and CI run this check.
+
+### 4. Update the schema diagram
+
+If you changed the schema in a meaningful way (added a table, added a foreign key, etc.), update the Mermaid diagram in `doc/code/memory/10_schema_diagram.md`.
+
+## How Migrations Run at Startup
+
+Schema migrations are triggered inside each memory class constructor (`SQLiteMemory.__init__` and `AzureSQLMemory.__init__`). When `skip_schema_migration=False` (the default), the inherited `_run_schema_migration()` method on `MemoryInterface` runs:
+
+```
+SQLiteMemory.__init__() / AzureSQLMemory.__init__()
+  → _run_schema_migration()                      # pyrit/memory/memory_interface.py
+      → run_schema_migrations(engine=...)         # pyrit/memory/migration.py
+          → alembic upgrade head
+      → check_schema_migrations(engine=...)       # pyrit/memory/migration.py
+          → alembic check
+```
+
+Both SQLite and AzureSQL follow the same migration path: first `run_schema_migrations` applies any pending Alembic revisions (`alembic upgrade head`), then `check_schema_migrations` verifies the resulting schema matches the current models (`alembic check`). The behavior depends on database state:
+
+| Database state | What happens |
+|---|---|
+| **Fresh (no tables)** | All migrations apply from scratch |
+| **Already versioned** | Only unapplied migrations run (idempotent) |
+| **Legacy (tables exist, no version tracking)** | Validates schema matches models, stamps current version, then upgrades. Raises `RuntimeError` on mismatch to prevent data corruption |
+
+Migrations run inside a transaction (`engine.begin()`), so a failed migration rolls back cleanly. The version tracking table is `pyrit_memory_alembic_version`.
+
+Users can skip migrations by passing `skip_schema_migration=True` to the memory class constructor. When using `initialize_pyrit_async()`, this can be forwarded via `**memory_instance_kwargs`:
+
+```python
+await initialize_pyrit_async("SQLite", skip_schema_migration=True)
+```
+
+## Important Rules
+
+### Migration revisions are immutable
+
+Once a migration revision is committed, it **must not be modified or deleted**. This is enforced by a pre-commit hook (`enforce_alembic_revision_immutability`). If you need to fix a migration, create a new revision instead.
+
+### Pre-commit hooks
+
+Two hooks run automatically when you touch memory-related files:
+
+1. **`enforce_alembic_revision_immutability`** — blocks modifications/deletions to existing revision files.
+2. **`memory-migrations-check`** — runs `memory_migrations.py check` to verify the schema is in sync.
+
+These hooks trigger on changes to `pyrit/memory/memory_models.py`, `pyrit/memory/migration.py`, and files under `pyrit/memory/alembic/`.

doc/getting_started/install_local.md

@@ -4,7 +4,7 @@ Install PyRIT directly into your Python environment for full control and easy in
 
 ## Prerequisites
 
-- Python 3.10, 3.11, 3.12, or 3.13 (check with `python --version`)
+- Python 3.10, 3.11, 3.12, 3.13, or 3.14 (check with `python --version`)
 
 ## Install with pip or uv
 

doc/myst.yml

@@ -53,6 +53,7 @@ project:
         - file: contributing/8_pre_commit.md
         - file: contributing/9_exception.md
         - file: contributing/10_release_process.md
+        - file: contributing/11_memory_models.md
     - file: gui/0_gui.md
     - file: scanner/0_scanner.md
       children:

doc/references.bib

@@ -519,6 +519,14 @@ @article{rottger2023xstest
   url       = {https://arxiv.org/abs/2308.01263},
 }
 
+@article{zong2024vlguard,
+  title     = {Safety Fine-Tuning at (Almost) No Cost: A Baseline for Vision Large Language Models},
+  author    = {Yongshuo Zong and Ondrej Bohdal and Tingyang Yu and Yongxin Yang and Timothy Hospedales},
+  journal   = {arXiv preprint arXiv:2402.02207},
+  year      = {2024},
+  url       = {https://arxiv.org/abs/2402.02207},
+}
+
 @article{vidgen2024ailuminate,
   title     = {Introducing v0.5 of the {AI} Safety Benchmark from {MLCommons}},
   author    = {Bertie Vidgen and Adarsh Agrawal and Ahmed M. Ahmed and Victor Akinwande and Namir Al-Nuaimi and others},