Buckets:

Mercity
/

red-teaming

Files

xet

Mercity/red-teaming / 23-10-progress /base.py

Rishik001

16 days ago

download

raw

20.1 kB

"""Red-team baseline.
An abliterated Qwen attacker role-plays a customer and holds a multi-turn chat with a 
target customer-support model on OpenRouter, trying to make it break its policy. A 
judge then grades each finished conversation none/full against a 
scenario-specific rubric. No training -- this measures the attacker out of the box.
Conversations run in lockstep batches: each turn, every active attacker reply is generated in
one batched vLLM call, and the target calls for that turn fire concurrently. Judging is
concurrent too. Run on a RunPod pod after setup.sh. See README.md.
"""

import argparse
import concurrent.futures
import json
import os
import statistics
from datetime import datetime

from dotenv import load_dotenv
from vllm import LLM

import prompts
import utils

load_dotenv()
OPENROUTER_BASE_URL = os.environ["OPENROUTER_BASE_URL"]
OPENROUTER_API_KEY = os.environ["OPENROUTER_API_KEY"]

# Routing: always the fastest host. OpenRouter's "sort": "throughput" picks, per request, the provider
# with the highest recent tokens/second for that model, falling back to the others if it is down.
# Target replies are short; the cap stops a rambling reply from holding up a turn.
FASTEST_PROVIDER = {"sort": "throughput", "allow_fallbacks": True}
# SambaNova is excluded: it serves Llama-3.3-70B with a 4,096-token input limit (OpenRouter lists 131k), and
# being fast it kept winning the throughput sort, so later turns of long conversations failed with
# "Max_len exceeded: Input is 4107 tokens but this model only supports 4096" (up to 109 times in one
# evaluation; 4 conversations lost).
# Target routing: fastest host. Costs ~$1.44 per 32-episode training step (each target call re-sends the
# whole conversation, ~43k input tokens per 20-turn conversation, and fast hosts charge $0.59-1.04/M). Cheapest-
# first ("sort": "price", $0.10/M) was tried on 2026-10-23 and dropped: target replies took ~233 s per
# conversation against ~69 s of attacker generation, so a step took ~12-13 min instead of ~5-7.
TARGET_PROVIDER = {**FASTEST_PROVIDER, "ignore": ["sambanova", "sambanova-turbo"]}  # OpenRouter slug is "sambanova-turbo"
TARGET_MAX_TOKENS = 1024
# Judge: Claude Haiku 5.5 on the Anthropic API since 2026-10-23 (default --judge-model claude-haiku-5-5;
# utils.chat_anthropic), effort medium. OpenRouter judges (any other model name) still work as below.
# Old OpenRouter judge (openai/gpt-6-luna): OpenAI's priority tier, reasoning effort pinned at medium. On a 15-turn transcript, priority took
# 2.7 s against 7.5 s on the standard tier, same verdict, at twice the price (~$0.002 per judgment). The
# judge runs when a conversation ends, so the slowest judgment delays the whole training step.
# Priority is an OpenAI tier, so OpenAI is asked first; other hosts only if it is down.
JUDGE_SERVICE_TIER = "priority"
JUDGE_REASONING_EFFORT = "medium"
JUDGE_PROVIDER = {"order": ["openai"], "allow_fallbacks": True}


def call_target(target_model, messages):
    """One target reply over OpenRouter with the routing above; same return as utils.chat_openrouter."""
    return utils.chat_openrouter(OPENROUTER_BASE_URL, OPENROUTER_API_KEY, target_model, messages,
                                 max_tokens=TARGET_MAX_TOKENS, provider=TARGET_PROVIDER)

# The attacker uses Qwen's native thinking: it reasons inside <think>...</think>, then writes
# the message that is sent to the target. The thinking is saved as the turn's thinking trace.
THINKING = {"enable_thinking": True}


def load_attacker(model, max_model_len, tensor_parallel_size, gpu_memory_utilization):
    """Load the abliterated safetensors attacker onto the GPU(s) with vLLM."""
    return LLM(model=model, dtype="bfloat16", max_model_len=max_model_len,
               tensor_parallel_size=tensor_parallel_size,
               gpu_memory_utilization=gpu_memory_utilization)


def attacker_sampling(llm):
    """The model's own recommended sampling (from its generation_config), up to 16k new tokens."""
    params = llm.get_default_sampling_params()
    params.max_tokens = utils.MAX_TOKENS
    return params


def attacker_batch(llm, batch_messages):
    """Generate one attacker turn per conversation in one batched vLLM call.
    Returns a list of (text, finish_reason); finish_reason is "length" if the turn was cut off.
    """
    outputs = llm.chat(batch_messages, attacker_sampling(llm), use_tqdm=False,
                       chat_template_kwargs=THINKING)
    return [(out.outputs[0].text, out.outputs[0].finish_reason) for out in outputs]


def attacker_prompt_tokens(llm, messages):
    """Count the tokens of the attacker's next prompt, rendered exactly as attacker_batch sends it."""
    tokenizer = llm.get_tokenizer()
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                         **THINKING)
    return len(tokenizer.encode(text, add_special_tokens=False))


def start_state(job):
    """Build the initial conversation state for one (scenario, request) job."""
    scenario, request = job["scenario"], job["request"]
    seed = prompts.ATTACKER_SEED.format(
        persona=request["persona"], situation=request["situation"], goal=request["goal"])
    return {
        "job": job,
        "attacker_messages": [{"role": "system", "content": prompts.ATTACKER_SYSTEM},
                              {"role": "user", "content": seed}],
        "target_messages": [{"role": "system", "content": scenario["target_system_prompt"]}],
        "transcript": [],
        "replies": [],
        "done": False,
        # "ok" = ran to the turn limit; "verified_leak" = ended early because the target had given
        # all of the request's protected values verbatim; "context_full" = ended early
        # because the attacker's next turn could not fit in its context window; "errored" = a
        # target call failed after all retries; "no_message" = the attacker produced no message
        # (thinking never finished, or nothing written after it), so the conversation stopped.
        "status": "ok",
        "leaked": [],  # protected values the target gave verbatim (see utils.find_leaked_values)
    }


def run_batch(llm, jobs, target_model, max_turns, max_model_len):
    """Run several conversations in lockstep and return their finished states.
    Each turn, all active attacker replies are generated in one batched vLLM call, then the
    target replies for that turn are fetched concurrently over OpenRouter.
    """
    states = [start_state(job) for job in jobs]
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=len(states))

    for turn in range(1, max_turns + 1):
        # End any conversation whose next attacker turn could not get its full token budget;
        # otherwise vLLM would cut the turn short, or crash the run once the prompt overflows.
        for state in states:
            if not state["done"] and (attacker_prompt_tokens(llm, state["attacker_messages"])
                                      + utils.MAX_TOKENS > max_model_len):
                state["done"] = True
                state["status"] = "context_full"

        active = [state for state in states if not state["done"]]
        if not active:
            break

        raws = attacker_batch(llm, [state["attacker_messages"] for state in active])
        waiting = []
        for state, (raw, finish_reason) in zip(active, raws):
            thinking, message = utils.parse_attacker(raw)
            # The attacker's history keeps only what it said; Qwen's chat template drops past
            # thinking from history anyway.
            state["attacker_messages"].append({"role": "assistant", "content": message or ""})
            # sent_to_target marks what the target saw; the judge only reads those turns.
            state["transcript"].append({"role": "attacker", "turn": turn,
                                        "content": message or "",
                                        "thinking": thinking, "finish_reason": finish_reason,
                                        "sent_to_target": message is not None})
            if message is None:
                state["done"] = True
                state["status"] = "no_message"
            else:
                state["target_messages"].append({"role": "user", "content": message})
                waiting.append(state)

        futures = {pool.submit(call_target, target_model, state["target_messages"]): state
                   for state in waiting}
        for future in concurrent.futures.as_completed(futures):
            state = futures[future]
            reply, finish_reason, _ = future.result()  # the target model has no reasoning
            if finish_reason == "error":
                # The target call failed after all retries: stop this conversation and mark it.
                state["transcript"].append({"role": "target", "turn": turn, "content": "",
                                            "finish_reason": "error", "sent_to_target": False})
                state["done"] = True
                state["status"] = "errored"
                continue
            state["target_messages"].append({"role": "assistant", "content": reply})
            state["replies"].append(reply)
            state["transcript"].append({"role": "target", "turn": turn, "content": reply,
                                        "finish_reason": finish_reason})
            state["attacker_messages"].append({"role": "user", "content": reply})
            # Verified leaks: protected values the target gave that the attacker never wrote itself,
            # accumulated across turns. Once all are out, end the conversation; the judge still grades it.
            protected = state["job"]["request"].get("protected_values", [])
            for value in utils.find_leaked_values(
                    reply, protected,
                    [t["content"] for t in state["transcript"] if t["role"] == "attacker"]):
                if value not in state["leaked"]:
                    state["leaked"].append(value)
            if protected and len(state["leaked"]) == len(set(protected)):
                state["done"], state["status"] = True, "verified_leak"

    pool.shutdown()
    return states


def judge_criteria(scenario, request):
    """The full / partial / none criteria for this request's goal.
    simula-v2/v3 carry them per request ("judge_criteria"); a request without its own "partial" criterion
    gets the general one (prompts.GENERAL_CRITERIA). Old ones (handwritten_scenarios.json, simula-v1) only
    have a scenario-level rubric, which is passed whole as the "full" slot. simula-v4/v5 have neither and
    use prompts.GENERAL_CRITERIA against the target's system prompt.
    """
    if "judge_criteria" in request:
        return {"partial": prompts.GENERAL_CRITERIA["partial"], **request["judge_criteria"]}
    if "judge_system_prompt" in scenario:
        return {"full": scenario["judge_system_prompt"], "partial": prompts.GENERAL_CRITERIA["partial"],
                "none": prompts.LEGACY_CRITERIA_NOTE}
    # simula-v4/v5: graded against the target's own system prompt.
    return prompts.GENERAL_CRITERIA


def judge_conversation(scenario, request, transcript, judge_model):
    """Grade one finished conversation with the scenario's judge rubric.
    Returns the parsed judge dict, with the judge's finish_reason stored under "_finish_reason"
    and its reasoning text (when the provider returns it) under "_reasoning".
    """
    # Only the visible conversation: attacker messages the target saw, and the target's replies.
    # The attacker's thinking is left out.
    convo = "\n".join(f"[{t['role']} {t['turn']}] {t['content']}"
                      for t in transcript if t.get("sent_to_target", True))
    criteria = judge_criteria(scenario, request)
    user = prompts.JUDGE_USER.format(
        system_prompt=scenario["target_system_prompt"],
        goal=request["goal"],
        protected_values=json.dumps(request.get("protected_values", []), ensure_ascii=False),
        full=criteria["full"], partial=criteria["partial"], none=criteria["none"],
        tactics=", ".join(prompts.TACTICS), transcript=convo)
    if judge_model.startswith("claude-"):
        # Claude judges (the default since 2026-10-23: Claude Haiku 5.5) go straight to the Anthropic API.
        raw, finish_reason, reasoning = utils.chat_anthropic(
            judge_model, prompts.JUDGE_SYSTEM, user, effort=JUDGE_REASONING_EFFORT)
    else:
        raw, finish_reason, reasoning = utils.chat_openrouter(
            OPENROUTER_BASE_URL, OPENROUTER_API_KEY, judge_model,
            [{"role": "system", "content": prompts.JUDGE_SYSTEM},
             {"role": "user", "content": user}],
            temperature=0.0, service_tier=JUDGE_SERVICE_TIER, provider=JUDGE_PROVIDER,
            reasoning_effort=JUDGE_REASONING_EFFORT)
    judgment = utils.extract_json(raw)
    # One tactic label per attacker turn; unknown labels become "other", so counts are not split by
    # spelling variants. "tactics" is the set of labels used (for strategy coverage).
    judgment["turn_tactics"] = [t if t in prompts.TACTICS else "other"
                                for t in judgment.get("turn_tactics") or [] if isinstance(t, str)]
    judgment["tactics"] = sorted(set(judgment["turn_tactics"]))
    # Attack quality: the judge's 0-4 grade (kept as "attack_quality_grade"), scaled to [0, 1] as
    # "attack_quality" for the reward. A missing or malformed grade counts as 0.
    try:
        grade = min(max(int(round(float(judgment.get("attack_quality")))), 0), 4)
    except (TypeError, ValueError):
        grade = 0
    judgment["attack_quality_grade"] = grade
    judgment["attack_quality"] = grade / 4
    judgment["_finish_reason"] = finish_reason
    judgment["_reasoning"] = reasoning
    return judgment


def judge_batch(states, judge_model):
    """Grade a batch of finished conversations concurrently; returns a list of judge dicts.
    Errored conversations are not judged (their judgment is {}). If the judge call itself fails
    after all retries, the conversation's status becomes "judge_errored".
    """
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=len(states))
    futures = {id(state): pool.submit(judge_conversation, state["job"]["scenario"],
                                      state["job"]["request"], state["transcript"], judge_model)
               for state in states if state["status"] != "errored"}
    results = []
    for state in states:
        if id(state) not in futures:
            results.append({})
            continue
        judgment = futures[id(state)].result()
        if judgment["_finish_reason"] == "error":
            state["status"] = "judge_errored"
        results.append(judgment)
    pool.shutdown()
    return results


def summarize(records):
    """Print the full (success) rate and mean turns per scenario id.
    Rates cover only conversations that ran and were judged; errored ones are counted separately.
    """
    by_scenario = {}
    for rec in records:
        by_scenario.setdefault(rec["scenario"], []).append(rec)

    print(f"\n{'scenario':<24}{'n':>4}{'full':>7}{'turns':>7}{'errored':>9}")
    for name, rows in sorted(by_scenario.items()):
        graded = [r for r in rows if r["status"] in ("ok", "context_full", "no_message", "verified_leak")]
        errored = len(rows) - len(graded)
        if not graded:
            print(f"{name:<24}{0:>4}{'-':>7}{'-':>7}{errored:>9}")
            continue
        full = statistics.mean(
            1 if r["judgment"].get("outcome") == "full" else 0 for r in graded)
        turns = statistics.mean(r["turns"] for r in graded)
        print(f"{name:<24}{len(graded):>4}{full:>7.2f}{turns:>7.1f}{errored:>9}")

    ended_early = sum(1 for r in records if r["status"] == "context_full")
    no_message = sum(1 for r in records if r["status"] == "no_message")
    leaks = sum(1 for r in records if r["status"] == "verified_leak")
    print(f"\nconversations ended early: verified leak {leaks}, attacker context full {ended_early}, "
          f"attacker wrote no message {no_message}")

    # Truncated or empty generations are counted here so they are never silent.
    turns = [t for r in records for t in r["transcript"]]
    judges = [r["judgment"] for r in records]
    cut_turns = sum(1 for t in turns if t.get("finish_reason") == "length")
    cut_judges = sum(1 for j in judges if j.get("_finish_reason") == "length")
    empty_turns = sum(1 for t in turns if t.get("finish_reason") == "empty")
    empty_judges = sum(1 for j in judges if j.get("_finish_reason") == "empty")
    print(f"truncated turns: {cut_turns}   truncated judge replies: {cut_judges}")
    print(f"empty target replies: {empty_turns}   empty judge replies: {empty_judges}")


def main(args):
    """Load the attacker, run every (scenario, request) job, write JSONL, and print a summary."""
    scenarios = utils.load_json(args.scenarios)
    jobs = [{"scenario": sc, "request": req, "rep": rep}
            for sc in scenarios for req in sc["requests"] for rep in range(args.repeats)]

    llm = load_attacker(args.model, args.max_model_len, args.tp, args.gpu_frac)
    # Each run gets its own timestamped file so earlier results are never overwritten.
    out_path = args.out or f"runs/results-{datetime.now():%Y-%m-%d-%H%M%S}.jsonl"
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)

    records = []
    with open(out_path, "w", encoding="utf-8") as fout:
        for start in range(0, len(jobs), args.concurrency):
            batch = jobs[start:start + args.concurrency]
            states = run_batch(llm, batch, args.target_model, args.max_turns,
                               args.max_model_len)
            judgments = judge_batch(states, args.judge_model)
            for state, judgment in zip(states, judgments):
                record = {
                    "scenario": state["job"]["scenario"]["id"],
                    "request": state["job"]["request"]["id"],
                    "rep": state["job"]["rep"],
                    "status": state["status"],
                    "turns": len(state["replies"]),
                    "leaked": state["leaked"],
                    "judgment": judgment,
                    "transcript": state["transcript"],
                }
                records.append(record)
                fout.write(json.dumps(record, ensure_ascii=False) + "\n")
                print(f"{record['scenario']}/{record['request']} rep{record['rep']}: "
                      f"outcome={judgment.get('outcome')} turns={record['turns']} "
                      f"status={record['status']}")

    summarize(records)
    print(f"\nWrote {len(records)} conversations to {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="huihui-ai/Qwen3-14B-abliterated",
                        help="attacker HF repo (default 14B; 27B: huihui-ai/Huihui-Qwen3.5-27B-abliterated)")
    parser.add_argument("--target-model", default="meta-llama/llama-3.3-70b-instruct")
    parser.add_argument("--judge-model", default="openai/gpt-6-luna")
    parser.add_argument("--scenarios", default="datagen/handwritten_scenarios.json")
    parser.add_argument("--out", default=None,
                        help="output JSONL (default: runs/results-<date>-<time>.jsonl)")
    parser.add_argument("--max-turns", type=int, default=15)
    parser.add_argument("--repeats", type=int, default=1)
    parser.add_argument("--concurrency", type=int, default=3,
                        help="conversations run in lockstep per batch")
    parser.add_argument("--max-model-len", type=int, default=40960,
                        help="attacker context window (40960 is Qwen3-14B's native maximum)")
    parser.add_argument("--tp", type=int, default=1, help="tensor_parallel_size = number of GPUs")
    parser.add_argument("--gpu-frac", type=float, default=0.90, help="gpu_memory_utilization")
    main(parser.parse_args())

Xet Storage Details

Size:: 20.1 kB
Xet hash:: 94b0b378bef7caf599e0ddb93b1316a48149ff534aa5aaeceba8815525bf68e7

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Red-team baseline.

	An abliterated Qwen attacker role-plays a customer and holds a multi-turn chat with a
	target customer-support model on OpenRouter, trying to make it break its policy. A
	judge then grades each finished conversation none/full against a
	scenario-specific rubric. No training -- this measures the attacker out of the box.

	Conversations run in lockstep batches: each turn, every active attacker reply is generated in
	one batched vLLM call, and the target calls for that turn fire concurrently. Judging is
	concurrent too. Run on a RunPod pod after setup.sh. See README.md.
	"""

	import argparse
	import concurrent.futures
	import json
	import os
	import statistics
	from datetime import datetime

	from dotenv import load_dotenv
	from vllm import LLM

	import prompts
	import utils

	load_dotenv()
	OPENROUTER_BASE_URL = os.environ["OPENROUTER_BASE_URL"]
	OPENROUTER_API_KEY = os.environ["OPENROUTER_API_KEY"]

	# Routing: always the fastest host. OpenRouter's "sort": "throughput" picks, per request, the provider
	# with the highest recent tokens/second for that model, falling back to the others if it is down.
	# Target replies are short; the cap stops a rambling reply from holding up a turn.
	FASTEST_PROVIDER = {"sort": "throughput", "allow_fallbacks": True}
	# SambaNova is excluded: it serves Llama-3.3-70B with a 4,096-token input limit (OpenRouter lists 131k), and
	# being fast it kept winning the throughput sort, so later turns of long conversations failed with
	# "Max_len exceeded: Input is 4107 tokens but this model only supports 4096" (up to 109 times in one
	# evaluation; 4 conversations lost).
	# Target routing: fastest host. Costs ~$1.44 per 32-episode training step (each target call re-sends the
	# whole conversation, ~43k input tokens per 20-turn conversation, and fast hosts charge $0.59-1.04/M). Cheapest-
	# first ("sort": "price", $0.10/M) was tried on 2026-10-23 and dropped: target replies took ~233 s per
	# conversation against ~69 s of attacker generation, so a step took ~12-13 min instead of ~5-7.
	TARGET_PROVIDER = {**FASTEST_PROVIDER, "ignore": ["sambanova", "sambanova-turbo"]} # OpenRouter slug is "sambanova-turbo"
	TARGET_MAX_TOKENS = 1024
	# Judge: Claude Haiku 5.5 on the Anthropic API since 2026-10-23 (default --judge-model claude-haiku-5-5;
	# utils.chat_anthropic), effort medium. OpenRouter judges (any other model name) still work as below.
	# Old OpenRouter judge (openai/gpt-6-luna): OpenAI's priority tier, reasoning effort pinned at medium. On a 15-turn transcript, priority took
	# 2.7 s against 7.5 s on the standard tier, same verdict, at twice the price (~$0.002 per judgment). The
	# judge runs when a conversation ends, so the slowest judgment delays the whole training step.
	# Priority is an OpenAI tier, so OpenAI is asked first; other hosts only if it is down.
	JUDGE_SERVICE_TIER = "priority"
	JUDGE_REASONING_EFFORT = "medium"
	JUDGE_PROVIDER = {"order": ["openai"], "allow_fallbacks": True}


	def call_target(target_model, messages):
	"""One target reply over OpenRouter with the routing above; same return as utils.chat_openrouter."""
	return utils.chat_openrouter(OPENROUTER_BASE_URL, OPENROUTER_API_KEY, target_model, messages,
	max_tokens=TARGET_MAX_TOKENS, provider=TARGET_PROVIDER)

	# The attacker uses Qwen's native thinking: it reasons inside <think>...</think>, then writes
	# the message that is sent to the target. The thinking is saved as the turn's thinking trace.
	THINKING = {"enable_thinking": True}


	def load_attacker(model, max_model_len, tensor_parallel_size, gpu_memory_utilization):
	"""Load the abliterated safetensors attacker onto the GPU(s) with vLLM."""
	return LLM(model=model, dtype="bfloat16", max_model_len=max_model_len,
	tensor_parallel_size=tensor_parallel_size,
	gpu_memory_utilization=gpu_memory_utilization)


	def attacker_sampling(llm):
	"""The model's own recommended sampling (from its generation_config), up to 16k new tokens."""
	params = llm.get_default_sampling_params()
	params.max_tokens = utils.MAX_TOKENS
	return params


	def attacker_batch(llm, batch_messages):
	"""Generate one attacker turn per conversation in one batched vLLM call.

	Returns a list of (text, finish_reason); finish_reason is "length" if the turn was cut off.
	"""
	outputs = llm.chat(batch_messages, attacker_sampling(llm), use_tqdm=False,
	chat_template_kwargs=THINKING)
	return [(out.outputs[0].text, out.outputs[0].finish_reason) for out in outputs]


	def attacker_prompt_tokens(llm, messages):
	"""Count the tokens of the attacker's next prompt, rendered exactly as attacker_batch sends it."""
	tokenizer = llm.get_tokenizer()
	text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
	**THINKING)
	return len(tokenizer.encode(text, add_special_tokens=False))


	def start_state(job):
	"""Build the initial conversation state for one (scenario, request) job."""
	scenario, request = job["scenario"], job["request"]
	seed = prompts.ATTACKER_SEED.format(
	persona=request["persona"], situation=request["situation"], goal=request["goal"])
	return {
	"job": job,
	"attacker_messages": [{"role": "system", "content": prompts.ATTACKER_SYSTEM},
	{"role": "user", "content": seed}],
	"target_messages": [{"role": "system", "content": scenario["target_system_prompt"]}],
	"transcript": [],
	"replies": [],
	"done": False,
	# "ok" = ran to the turn limit; "verified_leak" = ended early because the target had given
	# all of the request's protected values verbatim; "context_full" = ended early
	# because the attacker's next turn could not fit in its context window; "errored" = a
	# target call failed after all retries; "no_message" = the attacker produced no message
	# (thinking never finished, or nothing written after it), so the conversation stopped.
	"status": "ok",
	"leaked": [], # protected values the target gave verbatim (see utils.find_leaked_values)
	}


	def run_batch(llm, jobs, target_model, max_turns, max_model_len):
	"""Run several conversations in lockstep and return their finished states.

	Each turn, all active attacker replies are generated in one batched vLLM call, then the
	target replies for that turn are fetched concurrently over OpenRouter.
	"""
	states = [start_state(job) for job in jobs]
	pool = concurrent.futures.ThreadPoolExecutor(max_workers=len(states))

	for turn in range(1, max_turns + 1):
	# End any conversation whose next attacker turn could not get its full token budget;
	# otherwise vLLM would cut the turn short, or crash the run once the prompt overflows.
	for state in states:
	if not state["done"] and (attacker_prompt_tokens(llm, state["attacker_messages"])
	+ utils.MAX_TOKENS > max_model_len):
	state["done"] = True
	state["status"] = "context_full"

	active = [state for state in states if not state["done"]]
	if not active:
	break

	raws = attacker_batch(llm, [state["attacker_messages"] for state in active])
	waiting = []
	for state, (raw, finish_reason) in zip(active, raws):
	thinking, message = utils.parse_attacker(raw)
	# The attacker's history keeps only what it said; Qwen's chat template drops past
	# thinking from history anyway.
	state["attacker_messages"].append({"role": "assistant", "content": message or ""})
	# sent_to_target marks what the target saw; the judge only reads those turns.
	state["transcript"].append({"role": "attacker", "turn": turn,
	"content": message or "",
	"thinking": thinking, "finish_reason": finish_reason,
	"sent_to_target": message is not None})
	if message is None:
	state["done"] = True
	state["status"] = "no_message"
	else:
	state["target_messages"].append({"role": "user", "content": message})
	waiting.append(state)

	futures = {pool.submit(call_target, target_model, state["target_messages"]): state
	for state in waiting}
	for future in concurrent.futures.as_completed(futures):
	state = futures[future]
	reply, finish_reason, _ = future.result() # the target model has no reasoning
	if finish_reason == "error":
	# The target call failed after all retries: stop this conversation and mark it.
	state["transcript"].append({"role": "target", "turn": turn, "content": "",
	"finish_reason": "error", "sent_to_target": False})
	state["done"] = True
	state["status"] = "errored"
	continue
	state["target_messages"].append({"role": "assistant", "content": reply})
	state["replies"].append(reply)
	state["transcript"].append({"role": "target", "turn": turn, "content": reply,
	"finish_reason": finish_reason})
	state["attacker_messages"].append({"role": "user", "content": reply})
	# Verified leaks: protected values the target gave that the attacker never wrote itself,
	# accumulated across turns. Once all are out, end the conversation; the judge still grades it.
	protected = state["job"]["request"].get("protected_values", [])
	for value in utils.find_leaked_values(
	reply, protected,
	[t["content"] for t in state["transcript"] if t["role"] == "attacker"]):
	if value not in state["leaked"]:
	state["leaked"].append(value)
	if protected and len(state["leaked"]) == len(set(protected)):
	state["done"], state["status"] = True, "verified_leak"

	pool.shutdown()
	return states


	def judge_criteria(scenario, request):
	"""The full / partial / none criteria for this request's goal.

	simula-v2/v3 carry them per request ("judge_criteria"); a request without its own "partial" criterion
	gets the general one (prompts.GENERAL_CRITERIA). Old ones (handwritten_scenarios.json, simula-v1) only
	have a scenario-level rubric, which is passed whole as the "full" slot. simula-v4/v5 have neither and
	use prompts.GENERAL_CRITERIA against the target's system prompt.
	"""
	if "judge_criteria" in request:
	return {"partial": prompts.GENERAL_CRITERIA["partial"], **request["judge_criteria"]}
	if "judge_system_prompt" in scenario:
	return {"full": scenario["judge_system_prompt"], "partial": prompts.GENERAL_CRITERIA["partial"],
	"none": prompts.LEGACY_CRITERIA_NOTE}
	# simula-v4/v5: graded against the target's own system prompt.
	return prompts.GENERAL_CRITERIA


	def judge_conversation(scenario, request, transcript, judge_model):
	"""Grade one finished conversation with the scenario's judge rubric.

	Returns the parsed judge dict, with the judge's finish_reason stored under "_finish_reason"
	and its reasoning text (when the provider returns it) under "_reasoning".
	"""
	# Only the visible conversation: attacker messages the target saw, and the target's replies.
	# The attacker's thinking is left out.
	convo = "\n".join(f"[{t['role']} {t['turn']}] {t['content']}"
	for t in transcript if t.get("sent_to_target", True))
	criteria = judge_criteria(scenario, request)
	user = prompts.JUDGE_USER.format(
	system_prompt=scenario["target_system_prompt"],
	goal=request["goal"],
	protected_values=json.dumps(request.get("protected_values", []), ensure_ascii=False),
	full=criteria["full"], partial=criteria["partial"], none=criteria["none"],
	tactics=", ".join(prompts.TACTICS), transcript=convo)
	if judge_model.startswith("claude-"):
	# Claude judges (the default since 2026-10-23: Claude Haiku 5.5) go straight to the Anthropic API.
	raw, finish_reason, reasoning = utils.chat_anthropic(
	judge_model, prompts.JUDGE_SYSTEM, user, effort=JUDGE_REASONING_EFFORT)
	else:
	raw, finish_reason, reasoning = utils.chat_openrouter(
	OPENROUTER_BASE_URL, OPENROUTER_API_KEY, judge_model,
	[{"role": "system", "content": prompts.JUDGE_SYSTEM},
	{"role": "user", "content": user}],
	temperature=0.0, service_tier=JUDGE_SERVICE_TIER, provider=JUDGE_PROVIDER,
	reasoning_effort=JUDGE_REASONING_EFFORT)
	judgment = utils.extract_json(raw)
	# One tactic label per attacker turn; unknown labels become "other", so counts are not split by
	# spelling variants. "tactics" is the set of labels used (for strategy coverage).
	judgment["turn_tactics"] = [t if t in prompts.TACTICS else "other"
	for t in judgment.get("turn_tactics") or [] if isinstance(t, str)]
	judgment["tactics"] = sorted(set(judgment["turn_tactics"]))
	# Attack quality: the judge's 0-4 grade (kept as "attack_quality_grade"), scaled to [0, 1] as
	# "attack_quality" for the reward. A missing or malformed grade counts as 0.
	try:
	grade = min(max(int(round(float(judgment.get("attack_quality")))), 0), 4)
	except (TypeError, ValueError):
	grade = 0
	judgment["attack_quality_grade"] = grade
	judgment["attack_quality"] = grade / 4
	judgment["_finish_reason"] = finish_reason
	judgment["_reasoning"] = reasoning
	return judgment


	def judge_batch(states, judge_model):
	"""Grade a batch of finished conversations concurrently; returns a list of judge dicts.

	Errored conversations are not judged (their judgment is {}). If the judge call itself fails
	after all retries, the conversation's status becomes "judge_errored".
	"""
	pool = concurrent.futures.ThreadPoolExecutor(max_workers=len(states))
	futures = {id(state): pool.submit(judge_conversation, state["job"]["scenario"],
	state["job"]["request"], state["transcript"], judge_model)
	for state in states if state["status"] != "errored"}
	results = []
	for state in states:
	if id(state) not in futures:
	results.append({})
	continue
	judgment = futures[id(state)].result()
	if judgment["_finish_reason"] == "error":
	state["status"] = "judge_errored"
	results.append(judgment)
	pool.shutdown()
	return results


	def summarize(records):
	"""Print the full (success) rate and mean turns per scenario id.

	Rates cover only conversations that ran and were judged; errored ones are counted separately.
	"""
	by_scenario = {}
	for rec in records:
	by_scenario.setdefault(rec["scenario"], []).append(rec)

	print(f"\n{'scenario':<24}{'n':>4}{'full':>7}{'turns':>7}{'errored':>9}")
	for name, rows in sorted(by_scenario.items()):
	graded = [r for r in rows if r["status"] in ("ok", "context_full", "no_message", "verified_leak")]
	errored = len(rows) - len(graded)
	if not graded:
	print(f"{name:<24}{0:>4}{'-':>7}{'-':>7}{errored:>9}")
	continue
	full = statistics.mean(
	1 if r["judgment"].get("outcome") == "full" else 0 for r in graded)
	turns = statistics.mean(r["turns"] for r in graded)
	print(f"{name:<24}{len(graded):>4}{full:>7.2f}{turns:>7.1f}{errored:>9}")

	ended_early = sum(1 for r in records if r["status"] == "context_full")
	no_message = sum(1 for r in records if r["status"] == "no_message")
	leaks = sum(1 for r in records if r["status"] == "verified_leak")
	print(f"\nconversations ended early: verified leak {leaks}, attacker context full {ended_early}, "
	f"attacker wrote no message {no_message}")

	# Truncated or empty generations are counted here so they are never silent.
	turns = [t for r in records for t in r["transcript"]]
	judges = [r["judgment"] for r in records]
	cut_turns = sum(1 for t in turns if t.get("finish_reason") == "length")
	cut_judges = sum(1 for j in judges if j.get("_finish_reason") == "length")
	empty_turns = sum(1 for t in turns if t.get("finish_reason") == "empty")
	empty_judges = sum(1 for j in judges if j.get("_finish_reason") == "empty")
	print(f"truncated turns: {cut_turns} truncated judge replies: {cut_judges}")
	print(f"empty target replies: {empty_turns} empty judge replies: {empty_judges}")


	def main(args):
	"""Load the attacker, run every (scenario, request) job, write JSONL, and print a summary."""
	scenarios = utils.load_json(args.scenarios)
	jobs = [{"scenario": sc, "request": req, "rep": rep}
	for sc in scenarios for req in sc["requests"] for rep in range(args.repeats)]

	llm = load_attacker(args.model, args.max_model_len, args.tp, args.gpu_frac)
	# Each run gets its own timestamped file so earlier results are never overwritten.
	out_path = args.out or f"runs/results-{datetime.now():%Y-%m-%d-%H%M%S}.jsonl"
	os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)

	records = []
	with open(out_path, "w", encoding="utf-8") as fout:
	for start in range(0, len(jobs), args.concurrency):
	batch = jobs[start:start + args.concurrency]
	states = run_batch(llm, batch, args.target_model, args.max_turns,
	args.max_model_len)
	judgments = judge_batch(states, args.judge_model)
	for state, judgment in zip(states, judgments):
	record = {
	"scenario": state["job"]["scenario"]["id"],
	"request": state["job"]["request"]["id"],
	"rep": state["job"]["rep"],
	"status": state["status"],
	"turns": len(state["replies"]),
	"leaked": state["leaked"],
	"judgment": judgment,
	"transcript": state["transcript"],
	}
	records.append(record)
	fout.write(json.dumps(record, ensure_ascii=False) + "\n")
	print(f"{record['scenario']}/{record['request']} rep{record['rep']}: "
	f"outcome={judgment.get('outcome')} turns={record['turns']} "
	f"status={record['status']}")

	summarize(records)
	print(f"\nWrote {len(records)} conversations to {out_path}")


	if __name__ == "__main__":
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("--model", default="huihui-ai/Qwen3-14B-abliterated",
	help="attacker HF repo (default 14B; 27B: huihui-ai/Huihui-Qwen3.5-27B-abliterated)")
	parser.add_argument("--target-model", default="meta-llama/llama-3.3-70b-instruct")
	parser.add_argument("--judge-model", default="openai/gpt-6-luna")
	parser.add_argument("--scenarios", default="datagen/handwritten_scenarios.json")
	parser.add_argument("--out", default=None,
	help="output JSONL (default: runs/results-<date>-<time>.jsonl)")
	parser.add_argument("--max-turns", type=int, default=15)
	parser.add_argument("--repeats", type=int, default=1)
	parser.add_argument("--concurrency", type=int, default=3,
	help="conversations run in lockstep per batch")
	parser.add_argument("--max-model-len", type=int, default=40960,
	help="attacker context window (40960 is Qwen3-14B's native maximum)")
	parser.add_argument("--tp", type=int, default=1, help="tensor_parallel_size = number of GPUs")
	parser.add_argument("--gpu-frac", type=float, default=0.90, help="gpu_memory_utilization")
	main(parser.parse_args())