#!/usr/bin/env python3 """Build a conservative exploratory code-change plan.""" from __future__ import annotations import argparse import json import re from pathlib import Path from typing import Any, Dict, List SKIP_PARTS = { "__pycache__", ".git", "repro_outputs", "train_outputs", "analysis_outputs", "debug_outputs", "explore_outputs", "tmp", } CODE_SUFFIXES = {".py", ".yaml", ".yml", ".json", ".toml", ".ini"} MODEL_PATTERN = re.compile(r"(model|network|backbone|encoder|decoder|adapter|lora|head|loss)", re.IGNORECASE) TRAIN_PATTERN = re.compile(r"(train|trainer|optim|loss|config)", re.IGNORECASE) TASK_KEYWORDS = { "classification": ("class", "imagenet", "knn", "linear", "log_regression"), "segmentation": ("seg", "segment", "mask", "ade20k", "m2f", "mask2former"), "detection": ("det", "detect", "detr", "coco", "box"), "depth": ("depth", "nyu", "dpt", "depther"), "text": ("text", "token", "clip", "dinotxt"), "pretrain": ("pretrain", "ssl", "teacher", "student", "gram", "distillation"), } COMMON_TOKENS = {"py", "yaml", "yml", "json", "toml", "ini", "run", "train", "eval", "config", "configs"} def load_variant_spec(path: str) -> Dict[str, Any]: if not path: return {} return json.loads(Path(path).resolve().read_text(encoding="utf-8-sig")) def load_structured_payload(path: str) -> Any: if not path: return {} return json.loads(Path(path).resolve().read_text(encoding="utf-8-sig")) def normalize_task_family(value: Any) -> str: return str(value or "").strip().lower() def focus_tokens(current_research: str, task_family: str) -> List[str]: tokens: List[str] = [] for part in re.split(r"[^a-zA-Z0-9]+", current_research.lower()): if part and part not in COMMON_TOKENS and len(part) > 2: tokens.append(part) if task_family: tokens.append(task_family) tokens.extend(TASK_KEYWORDS.get(task_family, ())) ordered: List[str] = [] for token in tokens: if token not in ordered: ordered.append(token) return ordered[:20] def score_path(rel: str, task_family: str, tokens: List[str]) -> int: score = 0 if MODEL_PATTERN.search(rel): score += 5 if TRAIN_PATTERN.search(rel): score += 3 if rel.endswith(".py"): score += 1 lower = rel.lower() for token in TASK_KEYWORDS.get(task_family, ()): if token in lower: score += 4 for token in tokens: if token in lower: score += 2 if current_research_dir(rel, tokens): score += 3 return score def current_research_dir(rel: str, tokens: List[str]) -> bool: lower = rel.lower() slash_hits = [token for token in tokens if token in lower] return len(slash_hits) >= 2 def collect_candidate_edit_targets(repo: Path, current_research: str, task_family: str) -> List[str]: tokens = focus_tokens(current_research, task_family) scored: List[tuple[int, str]] = [] for path in repo.rglob("*"): if path.is_dir(): continue if any(part in SKIP_PARTS for part in path.relative_to(repo).parts): continue if path.suffix.lower() not in CODE_SUFFIXES: continue rel = path.relative_to(repo).as_posix() score = score_path(rel, task_family, tokens) if score: scored.append((score, rel)) scored.sort(key=lambda item: (-item[0], item[1])) return [rel for _, rel in scored[:8]] def select_idea_card(payload: Any) -> Dict[str, Any]: if isinstance(payload, dict): return payload if isinstance(payload, list) and payload: first = payload[0] if isinstance(first, dict): return first return {} def derive_target_location_map(targets: List[str], idea_card: Dict[str, Any], analysis: Dict[str, Any]) -> List[Dict[str, Any]]: config_hints = analysis.get("config_binding_hints", []) constructor_candidates = analysis.get("constructor_candidates", []) target_symbol = constructor_candidates[0] if constructor_candidates else (analysis.get("forward_candidates", []) or ["unspecified-symbol"])[0] results: List[Dict[str, Any]] = [] for path in targets[:4]: results.append( { "file": path, "role": "config" if path in config_hints else "code", "target_symbol": target_symbol, "reason": f"Maps `{idea_card.get('change_scope', 'candidate change')}` into `{idea_card.get('target_component', 'unspecified')}`.", } ) return results def derive_supporting_changes(spec: Dict[str, Any], idea_card: Dict[str, Any], analysis: Dict[str, Any]) -> List[str]: changes: List[str] = [] for item in idea_card.get("supporting_changes", []) or []: if item not in changes: changes.append(str(item)) for path in analysis.get("config_binding_hints", [])[:2]: changes.append(f"Review config binding in `{path}` for reversible wiring.") for axis in sorted((spec.get("variant_axes") or {}).keys())[:2]: changes.append(f"Keep `{axis}` plumbed through existing config or CLI surfaces.") unique: List[str] = [] for item in changes: if item not in unique: unique.append(item) return unique[:6] def derive_patch_surface_summary(target_location_map: List[Dict[str, Any]], supporting_changes: List[str]) -> Dict[str, Any]: code_targets = [item for item in target_location_map if item["role"] == "code"] config_targets = [item for item in target_location_map if item["role"] == "config"] surface_score = min(1.0, 0.15 + 0.10 * len(code_targets) + 0.05 * len(config_targets) + 0.04 * len(supporting_changes)) return { "surface_score": round(surface_score, 4), "code_target_count": len(code_targets), "config_target_count": len(config_targets), "summary": f"{len(code_targets)} code target(s), {len(config_targets)} config target(s), {len(supporting_changes)} supporting change(s).", } def derive_minimal_patch_plan( target_location_map: List[Dict[str, Any]], idea_card: Dict[str, Any], analysis: Dict[str, Any], ) -> List[Dict[str, Any]]: plan: List[Dict[str, Any]] = [] config_targets = [item["file"] for item in target_location_map if item["role"] == "config"] code_targets = [item["file"] for item in target_location_map if item["role"] == "code"] if config_targets: plan.append( { "change_type": "config-only", "target_files": config_targets, "rollback": "Revert the config override or remove the added config key.", "rationale": f"Expose `{idea_card.get('change_scope', 'candidate change')}` through frozen config surfaces first.", } ) if code_targets: plan.append( { "change_type": "import-glue", "target_files": [code_targets[0]], "rollback": "Remove the import/registry entry and restore the baseline route.", "rationale": "Keep wiring mechanical before any behavioral shim.", } ) plan.append( { "change_type": "module-transplant-shim", "target_files": [code_targets[0]], "rollback": "Delete the shim and return the call-site to the baseline symbol.", "rationale": "Only add a thin shim if constructor or forward surfaces do not already match.", } ) protected = analysis.get("metric_files", [])[:2] if protected: plan.append( { "change_type": "protected-zone-no-touch", "target_files": protected, "rollback": "No-op; evaluation and metric files should remain unchanged.", "rationale": "Preserve metric and leaderboard semantics unless the campaign explicitly allows mutation.", } ) return plan def derive_smoke_validation_plan( target_location_map: List[Dict[str, Any]], analysis: Dict[str, Any], spec: Dict[str, Any], ) -> List[Dict[str, Any]]: return [ { "name": "syntax-parse", "scope": [item["file"] for item in target_location_map if item["file"].endswith(".py")], "status": "planned", }, { "name": "import-resolution", "scope": [item["file"] for item in target_location_map if item["file"].endswith(".py")], "status": "planned", }, { "name": "config-path", "scope": [item["file"] for item in target_location_map if item["role"] == "config"], "status": "planned", }, { "name": "constructor-surface", "scope": analysis.get("constructor_candidates", [])[:4], "status": "planned", }, { "name": "forward-surface", "scope": analysis.get("forward_candidates", [])[:4], "status": "planned", }, { "name": "short-run-command", "scope": [str(spec.get("base_command") or "")], "status": "planned", }, ] def build_code_tracks(spec: Dict[str, Any], targets: List[str], task_family: str, current_research: str) -> List[str]: tracks: List[str] = [] if task_family: tracks.append(f"Stay anchored to the `{task_family}` task family while planning exploratory edits.") tracks.append(f"Preserve `{current_research}` as the comparison anchor for all code changes.") for axis, values in sorted((spec.get("variant_axes") or {}).items()): if not isinstance(values, (list, tuple)): values = [values] shown_values = ", ".join(str(value) for value in values[:3]) tracks.append(f"Review code touchpoints for `{axis}` variation across: {shown_values}.") if targets: tracks.append(f"Inspect candidate model files first: {', '.join(targets[:3])}.") if spec.get("base_command"): tracks.append(f"Keep `{spec['base_command']}` aligned with any exploratory code path changes.") tracks.extend( [ "Prefer one reversible module-level adaptation before broader rewrites.", "Keep config and entrypoint changes coupled so candidate runs remain attributable.", ] ) return tracks[:6] def build_payload( repo: Path, current_research: str, experiment_branch: str, spec: Dict[str, Any], task_family: str, idea_card: Dict[str, Any], analysis: Dict[str, Any], ) -> Dict[str, Any]: candidate_targets = collect_candidate_edit_targets(repo, current_research, task_family) target_location_map = derive_target_location_map(candidate_targets, idea_card, analysis) supporting_changes = derive_supporting_changes(spec, idea_card, analysis) patch_surface_summary = derive_patch_surface_summary(target_location_map, supporting_changes) minimal_patch_plan = derive_minimal_patch_plan(target_location_map, idea_card, analysis) smoke_validation_plan = derive_smoke_validation_plan(target_location_map, analysis, spec) code_tracks = build_code_tracks(spec, candidate_targets, task_family, current_research) return { "schema_version": "1.0", "repo": str(repo.resolve()), "current_research": current_research, "task_family": task_family or None, "experiment_branch": experiment_branch, "candidate_edit_targets": candidate_targets, "target_location_map": target_location_map, "supporting_changes": supporting_changes, "patch_surface_summary": patch_surface_summary, "minimal_patch_plan": minimal_patch_plan, "smoke_validation_plan": smoke_validation_plan, "proposed_code_tracks": code_tracks, "source_repo_refs": [ { "repo": repo.name, "ref": current_research, "note": "current_research anchor for exploratory code changes", } ], "notes": [ "Exploratory code plan only; candidate-level changes should stay isolated from the trusted baseline.", "Inspect candidate model files before introducing adapter or head changes.", ], } def main() -> int: parser = argparse.ArgumentParser(description="Build a conservative exploratory code-change plan.") parser.add_argument("--repo", required=True, help="Path to the target repository.") parser.add_argument("--current-research", required=True, help="Durable identifier for the current research context.") parser.add_argument("--experiment-branch", required=True, help="Isolated experiment branch label.") parser.add_argument("--variant-spec-json", default="", help="Optional path to the variant-spec JSON file.") parser.add_argument("--task-family", default="", help="Optional task-family hint used to focus candidate edit targets.") parser.add_argument("--idea-card-json", default="", help="Optional path to a selected idea-card JSON object or list.") parser.add_argument("--analysis-json", default="", help="Optional path to an analysis JSON object for richer structural hints.") parser.add_argument("--json", action="store_true", help="Emit JSON to stdout.") args = parser.parse_args() repo = Path(args.repo).resolve() idea_card = select_idea_card(load_structured_payload(args.idea_card_json)) analysis = load_structured_payload(args.analysis_json) payload = build_payload( repo, args.current_research, args.experiment_branch, load_variant_spec(args.variant_spec_json), normalize_task_family(args.task_family), idea_card, analysis if isinstance(analysis, dict) else {}, ) if args.json: print(json.dumps(payload, indent=2, ensure_ascii=False)) else: print(f"Current research: {payload['current_research']}") print(f"Task family: {payload.get('task_family') or 'unspecified'}") print(f"Experiment branch: {payload['experiment_branch']}") print("Candidate edit targets:", ", ".join(payload["candidate_edit_targets"]) or "none") print("Proposed code tracks:") for line in payload["proposed_code_tracks"]: print(f"- {line}") return 0 if __name__ == "__main__": raise SystemExit(main())