From 3a132a521a1d17751a1726ce2fb3c331aae3ac31 Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Wed, 15 Jul 2026 18:45:48 -0700 Subject: [PATCH] fix(ai-red-teaming): import get_generator in generated ATLAS script The ATLAS campaign template (`_build_atlas_imports`) omitted the `from dreadnode.generators.generator import get_generator, GenerateParams` import that the shared proxy-routing block relies on. Generated scripts raised `NameError: get_generator` at runtime (compile() only checks syntax, so it was not caught). Add the import and a static regression assertion. --- capabilities/ai-red-teaming/scripts/attack_runner.py | 1 + capabilities/ai-red-teaming/tests/test_attack_runner.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/capabilities/ai-red-teaming/scripts/attack_runner.py b/capabilities/ai-red-teaming/scripts/attack_runner.py index 20a1e50..7ede610 100644 --- a/capabilities/ai-red-teaming/scripts/attack_runner.py +++ b/capabilities/ai-red-teaming/scripts/attack_runner.py @@ -4701,6 +4701,7 @@ def _build_atlas_imports() -> str: "from pathlib import Path", "", "import dreadnode as dn", + "from dreadnode.generators.generator import get_generator, GenerateParams", "from dreadnode.airt.atlas import atlas_attack", "from dreadnode.airt.assessment import Assessment", "from dreadnode.airt.analytics import analyze", diff --git a/capabilities/ai-red-teaming/tests/test_attack_runner.py b/capabilities/ai-red-teaming/tests/test_attack_runner.py index b1fb9fd..c30f923 100644 --- a/capabilities/ai-red-teaming/tests/test_attack_runner.py +++ b/capabilities/ai-red-teaming/tests/test_attack_runner.py @@ -1006,6 +1006,14 @@ def test_generated_script_compiles_and_has_expected_pieces(self): assert "await atlas_attack(" in script assert 'SCENARIO_NAME = "finops"' in script assert "TOTAL_BUDGET = 8" in script + # Regression: the proxy-routing block calls get_generator(...), so the + # generated script MUST import it (compile() only checks syntax, not + # name resolution, so a missing import is a runtime NameError). + assert "get_generator(" in script + assert ( + "from dreadnode.generators.generator import get_generator, GenerateParams" + in script + ) # Default objective catalog is embedded with category coverage. assert "OBJECTIVES = " in script for cat in ("TW", "EA", "CB", "DE", "GH", "MP", "TB", "RP"):