#!/usr/bin/env python3 """Backfill reasoning for rows that lack it (rejection sampling / STaR). Consumes the ``*.to_synthesize.jsonl`` queue emitted by ``build_sft_dataset.py`` and produces ``*.synthesized.jsonl`` rows (with a verified ```` block) that ``build_sft_dataset.py --include-synthesized`` folds back into the ready mix. Two verification modes (carried on each record's ``verify`` block): label -- rejection sampling. Sample N teacher completions; keep the first whose final verdict matches the ground-truth label (vulnerable / not_vulnerable). This is the STaR / RFT idea applied to vuln detection: only correct reasoning survives. backfill -- the answer is fixed (secure-fix code, CVE write-up). Ask the teacher to produce a reasoning trace that justifies the given answer; accept a non-degenerate block and keep the original answer. Teacher endpoint is any OpenAI-compatible chat-completions server (e.g. a vLLM host). Configure via env or flags: TEACHER_BASE_URL (default http://localhost:8000/v1) TEACHER_MODEL TEACHER_API_KEY (default EMPTY) Use ``--mock`` to run the whole pipeline offline (no endpoint) for testing: it fabricates a deterministic, clearly-labelled trace so plumbing can be validated without a GPU. Mock output must never be used for a real Stage 1 run. """ from __future__ import annotations import argparse import json import os import sys import urllib.request from pathlib import Path from typing import Any THINK_INSTRUCTION = ( "Think step by step inside a single ... block, then give your " "final answer after it. Always include the block." ) def read_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as fh: for line in fh: line = line.strip() if line: rows.append(json.loads(line)) return rows def call_teacher(base_url: str, model: str, api_key: str, messages: list[dict[str, str]], temperature: float, max_tokens: int, timeout: int = 120) -> str: payload = { "model": model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, } req = urllib.request.Request( base_url.rstrip("/") + "/chat/completions", data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}, method="POST", ) with urllib.request.urlopen(req, timeout=timeout) as resp: body = json.loads(resp.read()) return body["choices"][0]["message"]["content"] def mock_completion(record: dict[str, Any], sample_idx: int) -> str: """Deterministic offline stand-in for the teacher (testing only).""" verify = record.get("verify", {}) or {} mode = verify.get("mode", "backfill") if mode == "label": expected = verify.get("expected", "not_vulnerable") cwe = ", ".join(verify.get("cwe", []) or []) or "no specific CWE" verdict = "Vulnerable." if expected == "vulnerable" else "Not vulnerable." return ( f"\n[mock trace] Inspecting the function for memory-safety and " f"input-validation issues; the evidence points to {expected} ({cwe}).\n\n\n" f"{verdict} Most likely {cwe}." ) answer = verify.get("answer", record.get("reference_answer", "")) return ( "\n[mock trace] Reasoning toward the provided secure answer: identify " "the flaw, then justify why the answer resolves it.\n\n\n" + answer ) def extract_think_and_body(text: str) -> tuple[str, str]: if "" in text and "" in text: start = text.index("") end = text.index("") + len("") return text[start:end], text[end:].strip() return "", text.strip() def verdict_of(text: str) -> str | None: low = text.lower() # check the negative phrase first so "not vulnerable" doesn't match "vulnerable" if "not vulnerable" in low or "no vulnerability" in low or "not a vulnerability" in low: return "not_vulnerable" if "vulnerable" in low or "vulnerability" in low: return "vulnerable" return None def think_is_degenerate(think: str) -> bool: inner = think.replace("", "").replace("", "").strip() return len(inner) < 40 def synthesize_one(record: dict[str, Any], args: argparse.Namespace) -> dict[str, Any] | None: verify = record.get("verify", {}) or {} mode = verify.get("mode", "backfill") prompt_messages = record.get("prompt_messages", []) teacher_messages = list(prompt_messages) if teacher_messages and teacher_messages[0].get("role") == "system": teacher_messages[0] = { "role": "system", "content": teacher_messages[0]["content"] + "\n\n" + THINK_INSTRUCTION, } else: teacher_messages.insert(0, {"role": "system", "content": THINK_INSTRUCTION}) for sample_idx in range(args.samples): if args.mock: completion = mock_completion(record, sample_idx) else: try: completion = call_teacher( args.base_url, args.model, args.api_key, teacher_messages, args.temperature, args.max_tokens, ) except Exception as exc: # noqa: BLE001 print(f"[warn] teacher call failed for {record.get('id')}: {exc!r}", file=sys.stderr) continue think, body = extract_think_and_body(completion) if not think or think_is_degenerate(think): continue if mode == "label": expected = verify.get("expected") if verdict_of(body) != expected: continue assistant = completion if "" in completion else f"{think}\n\n{body}" else: # backfill -> keep the known-good answer, prepend the verified reasoning answer = verify.get("answer", record.get("reference_answer", "")) assistant = f"{think}\n\n{answer}".strip() messages = list(prompt_messages) + [{"role": "assistant", "content": assistant}] return { "id": record.get("id"), "source": record.get("source"), "license": record.get("license", "missing"), "group": record.get("group"), "messages": messages, "metadata": { **(record.get("metadata") or {}), "think_status": "present", "synthesized": True, "synthesis_mode": mode, "synthesis_mock": bool(args.mock), }, } return None def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) parser.add_argument("--samples", type=int, default=4, help="Rollouts per row before giving up.") parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--max-tokens", type=int, default=1024) parser.add_argument("--limit", type=int, default=None, help="Only process the first N rows.") parser.add_argument("--base-url", default=os.getenv("TEACHER_BASE_URL", "http://localhost:8000/v1")) parser.add_argument("--model", default=os.getenv("TEACHER_MODEL", "teacher")) parser.add_argument("--api-key", default=os.getenv("TEACHER_API_KEY", "EMPTY")) parser.add_argument("--mock", action="store_true", help="Offline deterministic teacher (testing only).") return parser.parse_args() def main() -> int: args = parse_args() rows = read_jsonl(Path(args.input)) if args.limit is not None: rows = rows[: args.limit] out_path = Path(args.output) out_path.parent.mkdir(parents=True, exist_ok=True) kept = 0 failed = 0 with out_path.open("w", encoding="utf-8") as out: for record in rows: result = synthesize_one(record, args) if result is None: failed += 1 continue out.write(json.dumps(result, ensure_ascii=False, sort_keys=True) + "\n") kept += 1 print(json.dumps( { "input": args.input, "output": args.output, "rows_in": len(rows), "synthesized": kept, "failed": failed, "mock": bool(args.mock), "samples_per_row": args.samples, }, indent=2, )) return 0 if __name__ == "__main__": raise SystemExit(main())