From fc463a8a28ebed49ecad42b9292a1b90c769fe8a Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Wed, 1 Jul 2026 12:59:39 -0700 Subject: [PATCH 1/2] fix(ai-red-teaming): route dn/* models through proxy, keep local keys (ENG-7284) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generated attack workflow routed models through the LiteLLM proxy only when OPENAI_BASE_URL was set — but the platform (cloud sandbox and local TUI/CLI runtime) injects the proxy as DREADNODE_LLM_BASE / DREADNODE_LLM_API_KEY, so dn/* target/attacker/judge/transform models fell through to litellm unresolved and failed with "LLM Provider NOT provided". Resolve each model (target, attacker, judge, transform) with resolve_dn_model_to_generator: dn/* ids become proxy-configured generators via DREADNODE_LLM_*, while every other id (groq/, anthropic/, openai/, ...) is left untouched so litellm uses the user's own local provider API keys. The two modes now work side by side. Model-id strings are preserved for the platform UI labels. Verified end to end on dev across combinations (all-dn, dn+local mixes, all-local): trials run and route correctly with no provider errors. Bumps ai-red-teaming to 1.4.2. --- capabilities/ai-red-teaming/capability.yaml | 2 +- .../ai-red-teaming/scripts/attack_runner.py | 122 ++++++++---------- 2 files changed, 53 insertions(+), 71 deletions(-) diff --git a/capabilities/ai-red-teaming/capability.yaml b/capabilities/ai-red-teaming/capability.yaml index a8d093f..702c0df 100644 --- a/capabilities/ai-red-teaming/capability.yaml +++ b/capabilities/ai-red-teaming/capability.yaml @@ -1,6 +1,6 @@ schema: 1 name: ai-red-teaming -version: "1.4.1" +version: "1.4.2" description: > Probe the security and safety of AI applications, agents, and foundation models. Orchestrates adversarial attack workflows to discover vulnerabilities in LLMs, diff --git a/capabilities/ai-red-teaming/scripts/attack_runner.py b/capabilities/ai-red-teaming/scripts/attack_runner.py index 634c279..e125f24 100644 --- a/capabilities/ai-red-teaming/scripts/attack_runner.py +++ b/capabilities/ai-red-teaming/scripts/attack_runner.py @@ -795,7 +795,6 @@ def _auto_execute_workflow(filename: str, timeout: int = 3600) -> str: ATTACK_ALIASES["aprt"] = "aprt_progressive_attack" - _TRANSFORM_DEFS: dict[str, dict] = { # encoding "base64_encode": {"module": "dreadnode.transforms.encoding", "name": "base64_encode", "code": "base64_encode()"}, @@ -1026,21 +1025,21 @@ def _auto_execute_workflow(filename: str, timeout: int = 3600) -> str: "adapt_language": { "module": "dreadnode.transforms.language", "name": "adapt_language", - "code": 'adapt_language("Spanish", adapter_model=TRANSFORM_MODEL)', + "code": 'adapt_language("Spanish", adapter_model=TRANSFORM_MODEL_GEN)', "llm_powered": True, "parameterized": True, }, "code_switch": { "module": "dreadnode.transforms.language", "name": "code_switch", - "code": 'code_switch(["English", "Spanish"], adapter_model=TRANSFORM_MODEL, switch_ratio=0.4)', + "code": 'code_switch(["English", "Spanish"], adapter_model=TRANSFORM_MODEL_GEN, switch_ratio=0.4)', "llm_powered": True, "parameterized": True, }, "dialectal_variation": { "module": "dreadnode.transforms.language", "name": "dialectal_variation", - "code": 'dialectal_variation("AAVE", adapter_model=TRANSFORM_MODEL, intensity="moderate")', + "code": 'dialectal_variation("AAVE", adapter_model=TRANSFORM_MODEL_GEN, intensity="moderate")', "llm_powered": True, "parameterized": True, }, @@ -2721,7 +2720,7 @@ def _quote_arg_if_needed(arg: str) -> str: # Python identifier (e.g. TRANSFORM_MODEL, True, False, None) if re.match(r"^[A-Z_][A-Z_0-9]*$", arg) or arg in ("True", "False", "None"): return arg - # Keyword argument (e.g. adapter_model=TRANSFORM_MODEL) + # Keyword argument (e.g. adapter_model=TRANSFORM_MODEL_GEN) if "=" in arg: return arg # List literal @@ -2751,7 +2750,7 @@ def _resolve_transform(raw: str) -> dict: quoted_args = ", ".join(_quote_arg_if_needed(a) for a in _split_args(args_part)) code = "{}({})".format(defn["name"], quoted_args) if defn.get("llm_powered") and "adapter_model" not in args_part: - code = "{}({}, adapter_model=TRANSFORM_MODEL)".format(defn["name"], quoted_args) + code = "{}({}, adapter_model=TRANSFORM_MODEL_GEN)".format(defn["name"], quoted_args) return {**defn, "code": code, "resolved_name": canonical} key = raw.lower().replace("-", "_").replace(" ", "_") @@ -2821,6 +2820,7 @@ def _build_imports(attacks: list[dict], transforms: list[dict], has_scorers: boo "from dreadnode import task", "from dreadnode.generators.generator import get_generator, GenerateParams", "from dreadnode.generators.message import Message", + "from dreadnode.generators.proxy import resolve_dn_model_to_generator", ] attack_funcs = set() @@ -2995,69 +2995,50 @@ def _build_proxy_routing() -> str: a provider prefix. Models like groq/*, anthropic/*, together_ai/*, etc. are handled natively by litellm SDK using provider API keys from env. """ - return ''' -# Route LLM calls through the LiteLLM proxy when appropriate. -# Models with explicit provider prefixes (groq/, anthropic/, together_ai/, etc.) -# are routed directly by litellm SDK using provider API keys from the sandbox env. -# Only models without a provider prefix get routed through the proxy. + return """ +# Model routing — supports Dreadnode-proxied and bring-your-own-key models +# side by side, for the target, attacker, judge, and transform models alike: +# +# * dn/* -> routed through the platform LiteLLM proxy using the +# DREADNODE_LLM_BASE / DREADNODE_LLM_API_KEY that the platform injects +# (both the cloud sandbox and the local TUI/CLI runtime set these). +# * everything else (groq/*, anthropic/*, openai/*, together_ai/*, ...) is left +# untouched so litellm resolves it with the user's own local provider +# API keys (GROQ_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, ...). +# +# resolve_dn_model_to_generator() returns a proxy-configured Generator for dn/* +# ids and the original string for everything else. The MODEL string constants are +# preserved for the platform UI labels; the *_GEN values below drive inference. _DIRECT_PROVIDERS = ("groq/", "anthropic/", "together_ai/", "bedrock/", "azure/", "vertex_ai/", "cohere/", "replicate/", "mistral/", "ollama/", - "fireworks_ai/", "deepseek/", "huggingface/") -_proxy_key = os.environ.get("OPENAI_API_KEY", "") -_proxy_base = os.environ.get("OPENAI_BASE_URL", "") - -def _maybe_proxy(model_name: str) -> str: - """Prefix with openai/ only if model needs proxy routing.""" - if any(model_name.startswith(p) for p in _DIRECT_PROVIDERS): - return model_name # Direct provider — litellm SDK routes natively - if model_name.startswith("openai/"): - return model_name # Already prefixed - if _proxy_key and _proxy_base: - return f"openai/{model_name}" # Route through proxy - return model_name - -_orig_target = TARGET_MODEL -_orig_attacker = ATTACKER_MODEL -_orig_judge = JUDGE_MODEL -TARGET_MODEL = _maybe_proxy(TARGET_MODEL) -ATTACKER_MODEL = _maybe_proxy(ATTACKER_MODEL) -JUDGE_MODEL = _maybe_proxy(JUDGE_MODEL) -# Also proxy TRANSFORM_MODEL if it exists (used by LLM-powered transforms) + "fireworks_ai/", "deepseek/", "huggingface/", "openai/") + +TARGET_MODEL_GEN = resolve_dn_model_to_generator(TARGET_MODEL) +ATTACKER_MODEL_GEN = resolve_dn_model_to_generator(ATTACKER_MODEL) +JUDGE_MODEL_GEN = resolve_dn_model_to_generator(JUDGE_MODEL) try: - _orig_transform = TRANSFORM_MODEL - TRANSFORM_MODEL = _maybe_proxy(TRANSFORM_MODEL) + TRANSFORM_MODEL_GEN = resolve_dn_model_to_generator(TRANSFORM_MODEL) except NameError: - _orig_transform = None - -_any_proxied = (TARGET_MODEL != _orig_target or ATTACKER_MODEL != _orig_attacker - or JUDGE_MODEL != _orig_judge - or (_orig_transform is not None and TRANSFORM_MODEL != _orig_transform)) -if _any_proxied: - print(f" [proxy] Routing via {_proxy_base}") - if TARGET_MODEL != _orig_target: - print(f" [proxy] target: {_orig_target} -> {TARGET_MODEL}") - if ATTACKER_MODEL != _orig_attacker: - print(f" [proxy] attacker: {_orig_attacker} -> {ATTACKER_MODEL}") - if JUDGE_MODEL != _orig_judge: - print(f" [proxy] judge: {_orig_judge} -> {JUDGE_MODEL}") - if _orig_transform is not None and TRANSFORM_MODEL != _orig_transform: - print(f" [proxy] transform: {_orig_transform} -> {TRANSFORM_MODEL}") - sys.stdout.flush() + TRANSFORM_MODEL_GEN = None + +# The target task needs a concrete Generator; resolve_* returns a plain string for +# non-dn/ ids, so build one from it in that case. +TARGET_GENERATOR = ( + TARGET_MODEL_GEN if not isinstance(TARGET_MODEL_GEN, str) else get_generator(TARGET_MODEL_GEN) +) -# Warn if direct-provider models are missing their API key -_models_to_check = [("TARGET_MODEL", TARGET_MODEL, "target"), - ("ATTACKER_MODEL", ATTACKER_MODEL, "attacker"), - ("JUDGE_MODEL", JUDGE_MODEL, "judge")] -if _orig_transform is not None: - _models_to_check.append(("TRANSFORM_MODEL", TRANSFORM_MODEL, "transform")) -for _, _val, _label in _models_to_check: - if any(_val.startswith(p) for p in _DIRECT_PROVIDERS): - _provider = _val.split("/")[0].upper().replace("_AI", "") - _key_var = f"{_provider}_API_KEY" +# Report routing, and warn if a direct-provider model is missing its local key. +for _val, _label in ( + (TARGET_MODEL, "target"), (ATTACKER_MODEL, "attacker"), (JUDGE_MODEL, "judge") +): + if _val.startswith("dn/"): + print(f" [proxy] {_label}: {_val} -> Dreadnode LiteLLM proxy") + elif any(_val.startswith(p) for p in _DIRECT_PROVIDERS): + _key_var = _val.split("/")[0].upper().replace("_AI", "") + "_API_KEY" if not os.environ.get(_key_var): print(f" [warn] {_label} uses {_val} but {_key_var} not found in env") sys.stdout.flush() -''' +""" def _build_assessment_kwargs(config: dict, assessment_name: str, filename: str) -> str: @@ -3129,7 +3110,7 @@ def _build_target() -> str: return """\ @task async def target(prompt: str) -> str: - generator = get_generator(TARGET_MODEL) + generator = TARGET_GENERATOR messages = [Message(role="user", content=prompt)] last_error = None for attempt in range(3): @@ -3158,8 +3139,8 @@ def _build_attack_params( """Build the parameter string for an attack function call.""" params = ["goal={}".format(goal_expr), "target=target"] if atk["has_attacker"]: - params.append("attacker_model=ATTACKER_MODEL") - params.append("evaluator_model=JUDGE_MODEL") + params.append("attacker_model=ATTACKER_MODEL_GEN") + params.append("evaluator_model=JUDGE_MODEL_GEN") params.append("n_iterations=MAX_ITERATIONS") for k, v in atk.get("extra_defaults", {}).items(): params.append("{}={}".format(k, v)) @@ -3374,8 +3355,8 @@ def _generate_transform_study(config: dict) -> str: # Build attack params for the loop (transforms come from loop variable) params = ["goal=GOAL", "target=target"] if atk["has_attacker"]: - params.append("attacker_model=ATTACKER_MODEL") - params.append("evaluator_model=JUDGE_MODEL") + params.append("attacker_model=ATTACKER_MODEL_GEN") + params.append("evaluator_model=JUDGE_MODEL_GEN") params.append("n_iterations=MAX_ITERATIONS") for k, v in atk.get("extra_defaults", {}).items(): params.append("{}={}".format(k, v)) @@ -3732,8 +3713,8 @@ def _generate_category_attack(config: dict) -> str: # Build attack params for the template (goal comes from loop) params = ["goal=goal_text", "target=target"] if attacks[0]["has_attacker"]: - params.append("attacker_model=ATTACKER_MODEL") - params.append("evaluator_model=JUDGE_MODEL") + params.append("attacker_model=ATTACKER_MODEL_GEN") + params.append("evaluator_model=JUDGE_MODEL_GEN") params.append("n_iterations=MAX_ITERATIONS") for k, v in attacks[0].get("extra_defaults", {}).items(): params.append("{}={}".format(k, v)) @@ -3805,8 +3786,9 @@ def generate_category_attack(params: dict) -> dict: if not attack_names: return { "error": ( - "attacks must be one or more attack names, e.g. ['tap', 'goat'] " - "or 'tap,goat'. Got: {!r}".format(attacks_raw) + "attacks must be one or more attack names, e.g. ['tap', 'goat'] " "or 'tap,goat'. Got: {!r}".format( + attacks_raw + ) ) } From 365a4dacde053ea7c3e9cbfd9d35d56922235d70 Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Wed, 1 Jul 2026 13:48:10 -0700 Subject: [PATCH 2/2] fix(ai-red-teaming): inline dn/* resolver so workflow runs under any SDK build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generated workflow imported resolve_dn_model_to_generator from dreadnode.generators.proxy, but execute_workflow runs the workflow under the installed CLI tool's interpreter, whose proxy module can predate that symbol — producing "ImportError: cannot import name 'resolve_dn_model_to_generator'" and 0 trials. Inline the dn/* -> LiteLLM-proxy resolution (using only get_generator + GenerateParams, present in every SDK build) instead of importing it. Verified live: execute_workflow now succeeds on first try for dn/* (proxy) and non-dn/ (local provider key) models with no ImportError and no routing errors. --- .../ai-red-teaming/scripts/attack_runner.py | 36 ++++++++++++++----- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/capabilities/ai-red-teaming/scripts/attack_runner.py b/capabilities/ai-red-teaming/scripts/attack_runner.py index e125f24..252f8a7 100644 --- a/capabilities/ai-red-teaming/scripts/attack_runner.py +++ b/capabilities/ai-red-teaming/scripts/attack_runner.py @@ -2820,7 +2820,6 @@ def _build_imports(attacks: list[dict], transforms: list[dict], has_scorers: boo "from dreadnode import task", "from dreadnode.generators.generator import get_generator, GenerateParams", "from dreadnode.generators.message import Message", - "from dreadnode.generators.proxy import resolve_dn_model_to_generator", ] attack_funcs = set() @@ -3006,18 +3005,39 @@ def _build_proxy_routing() -> str: # untouched so litellm resolves it with the user's own local provider # API keys (GROQ_API_KEY, OPENAI_API_KEY, ANTHROPIC_API_KEY, ...). # -# resolve_dn_model_to_generator() returns a proxy-configured Generator for dn/* -# ids and the original string for everything else. The MODEL string constants are -# preserved for the platform UI labels; the *_GEN values below drive inference. +# _resolve_dn_model() returns a proxy-configured Generator for dn/* ids and the +# original string for everything else. It is inlined here (rather than imported +# from dreadnode.generators.proxy) so the workflow runs under ANY SDK build — +# including installed CLI tools whose proxy module predates that helper. The MODEL +# string constants are preserved for the platform UI labels; the *_GEN values below +# drive inference. _DIRECT_PROVIDERS = ("groq/", "anthropic/", "together_ai/", "bedrock/", "azure/", "vertex_ai/", "cohere/", "replicate/", "mistral/", "ollama/", "fireworks_ai/", "deepseek/", "huggingface/", "openai/") -TARGET_MODEL_GEN = resolve_dn_model_to_generator(TARGET_MODEL) -ATTACKER_MODEL_GEN = resolve_dn_model_to_generator(ATTACKER_MODEL) -JUDGE_MODEL_GEN = resolve_dn_model_to_generator(JUDGE_MODEL) +def _resolve_dn_model(model_name): + # dn/* -> LiteLLM-proxy Generator via DREADNODE_LLM_*; else the id unchanged. + if not isinstance(model_name, str) or not model_name.startswith("dn/"): + return model_name + _api_base = (os.environ.get("DREADNODE_LLM_BASE", "") or "").strip() or None + _api_key = (os.environ.get("DREADNODE_LLM_API_KEY", "") or "").strip() or None + if not _api_base or not _api_key: + raise RuntimeError( + "Missing proxy configuration — set DREADNODE_LLM_BASE and " + "DREADNODE_LLM_API_KEY to use " + model_name + ) + _gen = get_generator( + model_name, + params=GenerateParams(api_base=_api_base, extra={"custom_llm_provider": "litellm_proxy"}), + ) + _gen.api_key = _api_key + return _gen + +TARGET_MODEL_GEN = _resolve_dn_model(TARGET_MODEL) +ATTACKER_MODEL_GEN = _resolve_dn_model(ATTACKER_MODEL) +JUDGE_MODEL_GEN = _resolve_dn_model(JUDGE_MODEL) try: - TRANSFORM_MODEL_GEN = resolve_dn_model_to_generator(TRANSFORM_MODEL) + TRANSFORM_MODEL_GEN = _resolve_dn_model(TRANSFORM_MODEL) except NameError: TRANSFORM_MODEL_GEN = None