Class CodingAgent

nanoresearch/agents/coding.py:26–549

Generates runnable training code + SLURM scripts based on cloned repos and experiment plan.

class CodingAgent(_CodingHelpersMixin, BaseResearchAgent):
    """Generates runnable training code + SLURM scripts based on cloned repos and experiment plan."""

    stage = PipelineStage.CODING

    @property
    def stage_config(self):
        """Use code_gen model config for writing code."""
        return self.config.for_stage("code_gen")

    @staticmethod
    def _default_code_plan_files() -> list[dict[str, Any]]:
        return [
            {
                "path": "run_experiments.py",
                "description": "Unified entrypoint that loads configs/experiment_matrix.json, executes proposed/baseline/ablation/optimization/complexity runs, and writes all required result artifacts",
                "is_entrypoint": True,
            },
            {
                "path": "train.py",
                "description": "Single-run training/evaluation implementation used by run_experiments.py; supports --dry-run / --quick-eval",
                "is_entrypoint": False,
            },
            {"path": "model.py", "description": "Model and measured baseline definitions"},
            {"path": "dataset.py", "description": "Dataset loading and preprocessing"},
            {"path": "evaluate.py", "description": "Evaluation metrics and testing"},
            {"path": "config.py", "description": "Default hyperparameters and configuration"},
        ]

    def _normalize_code_plan(self, code_plan: dict[str, Any] | None) -> dict[str, Any]:
        plan = dict(code_plan) if isinstance(code_plan, dict) else {}

        normalized_files: list[dict[str, Any]] = []
        seen_paths: set[str] = set()
        for raw_spec in plan.get("files", []):
            if not isinstance(raw_spec, dict):
                continue
            path = str(raw_spec.get("path") or "").strip().replace("\\", "/")
            if not path or path in seen_paths:
                continue
            seen_paths.add(path)
            normalized_files.append(
                {
                    "path": path,
                    "description": str(raw_spec.get("description") or "").strip(),
                    "is_entrypoint": bool(raw_spec.get("is_entrypoint", False)),
                }
            )

        for default_spec in self._default_code_plan_files():
            if default_spec["path"] in seen_paths:
                continue
            normalized_files.append(dict(default_spec))
            seen_paths.add(default_spec["path"])

        dependencies: list[str] = []
        seen_dependencies: set[str] = set()
        for raw_dependency in plan.get("dependencies", []):