From 856802383916ae7f23e2d117a8c9ca934e75357a Mon Sep 17 00:00:00 2001 From: Kai Xu Date: Wed, 12 Aug 2026 22:33:20 +0000 Subject: [PATCH] feat: replace OptimizationLoop with optimize workflow graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add optimize_workflow() as a 7-node DAG in the workflow engine, unifying the optimization pipeline with all other factory modes. - Add workflow graph: baseline → gate_baseline → mutate → apply → execute → gate_improve (RELOOP max 5) → test_eval - Add `factory optimize-step` CLI with 4 subcommands (run-dev, run-test, apply-patch, check-gate) as thin wrappers around existing HarborBenchmark, evaluators, and evaluate_gate - Wire `factory optimize` to call `factory workflow run optimize` with --legacy flag for backward compat via OptimizationLoop - Add structlog deprecation warning to OptimizationLoop.__init__() - Add tests for workflow graph structure, optimize-step logic, and --legacy CLI routing Co-Authored-By: Claude Opus 4.6 --- factory/cli/__init__.py | 6 + factory/cli/_main.py | 13 +- factory/cli/_parser_groups.py | 15 ++ factory/cli/optimize.py | 111 ++++++++++++- factory/cli/optimize_step.py | 281 ++++++++++++++++++++++++++++++++ factory/optimization/loop.py | 3 + factory/workflow/definitions.py | 93 +++++++++++ tests/test_cli_optimize.py | 66 +++++++- tests/test_optimize_step.py | 237 +++++++++++++++++++++++++++ tests/test_workflow_optimize.py | 214 ++++++++++++++++++++++++ 10 files changed, 1027 insertions(+), 12 deletions(-) create mode 100644 factory/cli/optimize_step.py create mode 100644 tests/test_optimize_step.py create mode 100644 tests/test_workflow_optimize.py diff --git a/factory/cli/__init__.py b/factory/cli/__init__.py index cee6d4c0c..0bcb11db9 100644 --- a/factory/cli/__init__.py +++ b/factory/cli/__init__.py @@ -48,6 +48,12 @@ from factory.cli.optimize import ( cmd_optimize as cmd_optimize, ) +from factory.cli.optimize_step import ( + cmd_optimize_step_apply_patch as cmd_optimize_step_apply_patch, + cmd_optimize_step_check_gate as cmd_optimize_step_check_gate, + cmd_optimize_step_run_dev as cmd_optimize_step_run_dev, + cmd_optimize_step_run_test as cmd_optimize_step_run_test, +) from factory.cli.skillopt import ( cmd_skillopt as cmd_skillopt, ) diff --git a/factory/cli/_main.py b/factory/cli/_main.py index 4058f3a8b..5ae59a20d 100644 --- a/factory/cli/_main.py +++ b/factory/cli/_main.py @@ -105,7 +105,7 @@ "backfill-archive", ], ), - ("Self-Evolution", ["ace", "ace-stats", "digest", "workflow", "graph", "skillopt", "optimize"]), + ("Self-Evolution", ["ace", "ace-stats", "digest", "workflow", "graph", "skillopt", "optimize", "optimize-step"]), ( "Configuration", [ @@ -301,6 +301,17 @@ def main(argv: list[str] | None = None) -> int: "agent": _cli.cmd_agent, "skillopt": _cli.cmd_skillopt, "optimize": _cli.cmd_optimize, + "optimize-step": lambda a: { + "run-dev": _cli.cmd_optimize_step_run_dev, + "run-test": _cli.cmd_optimize_step_run_test, + "apply-patch": _cli.cmd_optimize_step_apply_patch, + "check-gate": _cli.cmd_optimize_step_check_gate, + }.get( + str(getattr(a, "optimize_step_command", "")), + lambda args: ( + print("Usage: factory optimize-step {run-dev,run-test,apply-patch,check-gate}") or 1 + ), + )(a), "ceo": _cli.cmd_ceo, "run": _cli.cmd_run, "tmux": _cli.cmd_tmux, diff --git a/factory/cli/_parser_groups.py b/factory/cli/_parser_groups.py index 26846ed61..f60d21e4a 100644 --- a/factory/cli/_parser_groups.py +++ b/factory/cli/_parser_groups.py @@ -253,6 +253,19 @@ def add_self_evolution_parsers(sub: argparse._SubParsersAction) -> None: # type p.add_argument("--epochs", type=int, default=1, help="Number of training epochs") p.add_argument("--steps-per-epoch", type=int, default=1, help="Steps per epoch") + opt_step_parser = sub.add_parser("optimize-step", help="Workflow node helpers for optimize graph") + opt_step_sub = opt_step_parser.add_subparsers(dest="optimize_step_command") + p_rd = opt_step_sub.add_parser("run-dev", help="Run benchmark dev split") + p_rd.add_argument("--project", required=True, help="Path to the project") + p_rt = opt_step_sub.add_parser("run-test", help="Run benchmark test split") + p_rt.add_argument("--project", required=True, help="Path to the project") + p_ap = opt_step_sub.add_parser("apply-patch", help="Apply mutation rules to current skill") + p_ap.add_argument("--project", required=True, help="Path to the project") + p_cg = opt_step_sub.add_parser("check-gate", help="Check optimization gate verdict") + p_cg.add_argument("--project", required=True, help="Path to the project") + p_cg.add_argument("--baseline", action="store_true", default=False, + help="Check baseline score instead of improvement gate") + p = sub.add_parser("optimize", help="Run inner-outer optimization loop with HarborBenchmark") p.add_argument("path", help="Path to the project") p.add_argument("--benchmark", default="searchqa", choices=["searchqa", "featurebench", "auto"], @@ -268,6 +281,8 @@ def add_self_evolution_parsers(sub: argparse._SubParsersAction) -> None: # type p.add_argument("--model", default="sonnet", help="Model for AgenticMutator (default: sonnet)") p.add_argument("--split-seed", type=int, default=42, help="Seed for reproducible split generation (default: 42)") p.add_argument("--splits-dir", default=None, help="Path to pre-generated JSONL split files") + p.add_argument("--legacy", action="store_true", default=False, + help="Use legacy OptimizationLoop instead of workflow graph") def add_configuration_parsers(sub: argparse._SubParsersAction) -> None: # type: ignore[type-arg] diff --git a/factory/cli/optimize.py b/factory/cli/optimize.py index 18eb02769..dd6655c75 100644 --- a/factory/cli/optimize.py +++ b/factory/cli/optimize.py @@ -3,18 +3,20 @@ from __future__ import annotations import argparse +import json import os +import subprocess import sys from pathlib import Path +from typing import Any +import structlog + +log = structlog.get_logger() -def cmd_optimize(args: argparse.Namespace) -> int: - """Run the optimization loop with HarborBenchmark executor.""" - project = Path(args.path).resolve() - if not project.exists(): - print(f"Error: project path does not exist: {project}", file=sys.stderr) - return 1 +def _run_legacy(args: argparse.Namespace, project: Path) -> int: + """Legacy path: use OptimizationLoop directly.""" from factory.optimization import AgenticMutator, LoopConfig, OptimizationLoop, Surface from factory.optimization.benchmarks.harbor import HarborBenchmark from factory.optimization.protocols import Evaluator, Executor @@ -170,3 +172,100 @@ def cmd_optimize(args: argparse.Namespace) -> int: print(f"Total improvement: {total_delta:+.4f}") return 0 + + +def cmd_optimize(args: argparse.Namespace) -> int: + """Run the optimization loop with HarborBenchmark executor.""" + project = Path(args.path).resolve() + if not project.exists(): + print(f"Error: project path does not exist: {project}", file=sys.stderr) + return 1 + + legacy = getattr(args, "legacy", False) + if legacy: + return _run_legacy(args, project) + + benchmark = getattr(args, "benchmark", None) or "searchqa" + git_ref = getattr(args, "git_ref", None) or os.environ.get("FACTORY_GIT_REF", "main") + docker_host = getattr(args, "docker_host", None) or os.environ.get("DOCKER_HOST", "") + concurrency = getattr(args, "concurrency", 5) + steps = getattr(args, "steps", 3) + model = getattr(args, "model", None) or "sonnet" + skill_path = getattr(args, "skill_path", None) + + # Setup: write initial state files + opt_dir = project / ".factory" / "optimization" + opt_dir.mkdir(parents=True, exist_ok=True) + + # Initial skill + skill_file = opt_dir / "current_skill.md" + if skill_path: + sp = Path(skill_path) + if sp.exists(): + skill_file.write_text(sp.read_text()) + elif not skill_file.exists(): + default_skill = ( + "# Question Answering Skill\n\n" + "(No learned rules yet.)\n\n" + "## Instructions\n\n" + "Read the question and search results from /tmp/task-instruction.md.\n" + "Answer the question and write ONLY your final answer to /workspace/answer.txt.\n" + "Also include your answer in tags in your response.\n" + ) + skill_file.write_text(default_skill) + + # Initial state.json + state_file = opt_dir / "state.json" + initial_state = { + "step": 0, + "current_score": 0.0, + "best_score": 0.0, + "best_step": 0, + "history": [], + } + state_file.write_text(json.dumps(initial_state, indent=2) + "\n") + + # Pass config to optimize-step via env vars + env = os.environ.copy() + env["FACTORY_OPT_BENCHMARK"] = benchmark + env["FACTORY_OPT_CONCURRENCY"] = str(concurrency) + env["FACTORY_GIT_REF"] = git_ref + env["FACTORY_OPT_MODEL"] = model + env["FACTORY_OPT_MAX_ITERATIONS"] = str(steps) + if docker_host: + env["DOCKER_HOST"] = docker_host + + print(f"Starting optimization (workflow): benchmark={benchmark}, max_iterations={steps}, concurrency={concurrency}") + + # Execute workflow + cmd = ["factory", "workflow", "run", "optimize", str(project)] + result = subprocess.run(cmd, env=env) + + # Read results + test_result_path = opt_dir / "test_result.json" + state: dict[str, Any] = json.loads(state_file.read_text()) if state_file.exists() else dict(initial_state) + + print(f"\n{'='*50}") + history: list[dict[str, Any]] = state.get("history", []) + print(f"Training complete: {len(history)} steps") + + if history: + baseline_score: float = history[0].get("score_start", 0.0) + print(f"Baseline score: {baseline_score:.4f}") + for h in history: + delta: float = h.get("score_delta", 0.0) + print(f" Step {h['step']}: {h['score_start']:.4f} -> {h['score_end']:.4f} " + f"({delta:+.4f}) verdict={h.get('verdict', 'n/a')}") + + print(f"Best score: {state['best_score']:.4f} (step {state['best_step']})") + print(f"Final score: {state['current_score']:.4f}") + + if test_result_path.exists(): + test_data = json.loads(test_result_path.read_text()) + print(f"Test score: {test_data['score']:.4f}") + + if history: + total_delta_f: float = state["current_score"] - history[0].get("score_end", 0.0) + print(f"Total improvement: {total_delta_f:+.4f}") + + return result.returncode diff --git a/factory/cli/optimize_step.py b/factory/cli/optimize_step.py new file mode 100644 index 000000000..ed183a26c --- /dev/null +++ b/factory/cli/optimize_step.py @@ -0,0 +1,281 @@ +"""CLI handler for factory optimize-step — thin wrappers for workflow graph nodes.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from pathlib import Path +from typing import Any + +import structlog + +log = structlog.get_logger() + + +def _opt_dir(project: Path) -> Path: + return project / ".factory" / "optimization" + + +def _read_state(project: Path) -> dict: + state_path = _opt_dir(project) / "state.json" + if state_path.exists(): + return json.loads(state_path.read_text()) + return {"step": 0, "current_score": 0.0, "best_score": 0.0, "best_step": 0, "history": []} + + +def _write_state(project: Path, state: dict) -> None: + state_path = _opt_dir(project) / "state.json" + state_path.parent.mkdir(parents=True, exist_ok=True) + state_path.write_text(json.dumps(state, indent=2) + "\n") + + +def _get_benchmark_and_evaluator( + project: Path, +) -> tuple[Any, Any, Any]: + """Build executor + evaluator from env vars, reusing existing protocol implementations.""" + benchmark = os.environ.get("FACTORY_OPT_BENCHMARK", "searchqa") + concurrency = int(os.environ.get("FACTORY_OPT_CONCURRENCY", "5")) + git_ref = os.environ.get("FACTORY_GIT_REF", "main") + docker_host = os.environ.get("DOCKER_HOST", "") + model = os.environ.get("FACTORY_OPT_MODEL", "sonnet") + + from factory.optimization.benchmarks.harbor import HarborBenchmark + from factory.optimization.surface import Surface + from factory.optimization.types import BenchmarkSplits + + splits: BenchmarkSplits | None = None + splits_dir = project / ".factory" / "eval" / "benchmark" / "splits" + if splits_dir.is_dir(): + splits = BenchmarkSplits.from_jsonl_dir(splits_dir) + + executor: Any + evaluator: Any + + match benchmark: + case "searchqa": + from factory.optimization.benchmarks.searchqa import SearchQAEvaluator + + executor = HarborBenchmark( + git_ref=git_ref, concurrency=concurrency, + docker_host=docker_host, model=model, splits=splits, + ) + evaluator = SearchQAEvaluator() + + case "featurebench": + from factory.optimization.benchmarks.featurebench import FeatureBenchEvaluator + + executor = HarborBenchmark( + git_ref=git_ref, concurrency=concurrency, + docker_host=docker_host, model=model, splits=splits, + dataset="featurebench", + agent_class="factory_harbor_agent:FeatureBenchFactoryCeo", + ) + evaluator = FeatureBenchEvaluator() + + case _: + from factory.optimization.benchmarks.loader import load_benchmark + + benchmark_dir = project / ".factory" / "eval" / "benchmark" + defn = load_benchmark(benchmark_dir) + executor_params = defn.config.get("executor_params", {}) + evaluator_params = defn.config.get("evaluator_params", {}) + executor = defn.executor_cls(**executor_params) + evaluator = defn.evaluator_cls(**evaluator_params) + + skill_path = _opt_dir(project) / "current_skill.md" + skill_text = skill_path.read_text() if skill_path.exists() else "" + surface = Surface(prompt_slots={"skill": skill_text}) + + return executor, evaluator, surface + + +def cmd_optimize_step_run_dev(args: argparse.Namespace) -> int: + """Run the dev split of the benchmark and record results.""" + project = Path(args.project).resolve() + executor, evaluator, surface = _get_benchmark_and_evaluator(project) + + execution_result = executor.execute(project, surface, split="dev") + + score = 0.0 + artifacts = [Path(a) for a in execution_result.artifacts] + if artifacts: + eval_result = evaluator.parse_many(artifacts) + if eval_result.valid: + score = eval_result.score + + state = _read_state(project) + step = state["step"] + 1 + state["step"] = step + + step_record = { + "step": step, + "score_start": state["current_score"], + "score_end": score, + "score_delta": score - state["current_score"], + "verdict": "pending", + } + state["history"].append(step_record) + state["current_score"] = score + + if score > state["best_score"]: + state["best_score"] = score + state["best_step"] = step + + _write_state(project, state) + + # Write per-step artifacts + step_dir = _opt_dir(project) / "steps" / str(step) + step_dir.mkdir(parents=True, exist_ok=True) + + results_data = [] + if hasattr(execution_result, "task_results") and execution_result.task_results: + results_data = [ + {"task_id": t.task_id, "reward": t.reward, "predicted": t.predicted, "gold": t.gold} + for t in execution_result.task_results + ] + (step_dir / "results.json").write_text(json.dumps(results_data, indent=2) + "\n") + + skill_path = _opt_dir(project) / "current_skill.md" + if skill_path.exists(): + import shutil + shutil.copy2(skill_path, step_dir / "skill.md") + + # On first call (step=1), also write baseline.json + if step == 1: + baseline = {"score": score, "step": 1, "task_results": results_data} + (_opt_dir(project) / "baseline.json").write_text(json.dumps(baseline, indent=2) + "\n") + + log.info("optimize_step.run_dev", step=step, score=round(score, 4)) + print(json.dumps({"step": step, "score": score})) + return 0 + + +def cmd_optimize_step_run_test(args: argparse.Namespace) -> int: + """Run the test split for a final unbiased score.""" + project = Path(args.project).resolve() + executor, evaluator, surface = _get_benchmark_and_evaluator(project) + + execution_result = executor.execute(project, surface, split="test") + + score = 0.0 + artifacts = [Path(a) for a in execution_result.artifacts] + if artifacts: + eval_result = evaluator.parse_many(artifacts) + if eval_result.valid: + score = eval_result.score + + results_data = [] + if hasattr(execution_result, "task_results") and execution_result.task_results: + results_data = [ + {"task_id": t.task_id, "reward": t.reward, "predicted": t.predicted, "gold": t.gold} + for t in execution_result.task_results + ] + + test_result = {"score": score, "task_results": results_data} + result_path = _opt_dir(project) / "test_result.json" + result_path.write_text(json.dumps(test_result, indent=2) + "\n") + + log.info("optimize_step.run_test", score=round(score, 4)) + print(json.dumps({"score": score})) + return 0 + + +def cmd_optimize_step_apply_patch(args: argparse.Namespace) -> int: + """Read mutation.json, append rules to current_skill.md.""" + project = Path(args.project).resolve() + mutation_path = _opt_dir(project) / "mutation.json" + + if not mutation_path.exists(): + print("Error: mutation.json not found", file=sys.stderr) + return 1 + + raw = mutation_path.read_text().strip() + + # Parse JSON — with regex fallback for markdown-wrapped JSON + try: + data = json.loads(raw) + except json.JSONDecodeError: + match = re.search(r"\{[^{}]*\"rules\"[^{}]*\}", raw, re.DOTALL) + if match: + try: + data = json.loads(match.group()) + except json.JSONDecodeError: + print("Error: could not parse mutation.json", file=sys.stderr) + return 1 + else: + print("Error: could not parse mutation.json", file=sys.stderr) + return 1 + + rules = data.get("rules", []) + if not rules: + log.info("optimize_step.apply_patch.no_rules") + return 0 + + skill_path = _opt_dir(project) / "current_skill.md" + skill_text = skill_path.read_text() if skill_path.exists() else "" + + new_rules = "\n## Learned Rules\n\n" + for rule in rules: + new_rules += f"- {rule}\n" + + skill_text += new_rules + skill_path.write_text(skill_text) + + log.info("optimize_step.apply_patch", n_rules=len(rules)) + print(json.dumps({"rules_applied": len(rules)})) + return 0 + + +def cmd_optimize_step_check_gate(args: argparse.Namespace) -> int: + """Check gate: PROCEED (exit 0), RELOOP (exit 1), HALT (exit 2).""" + project = Path(args.project).resolve() + is_baseline = getattr(args, "baseline", False) + + if is_baseline: + baseline_path = _opt_dir(project) / "baseline.json" + if not baseline_path.exists(): + state = _read_state(project) + score = state.get("current_score", 0.0) + else: + data = json.loads(baseline_path.read_text()) + score = data.get("score", 0.0) + if score > 0: + print("PROCEED") + return 0 + else: + print("HALT: baseline score is 0") + return 2 + + state = _read_state(project) + history = state.get("history", []) + max_iterations = int(os.environ.get("FACTORY_OPT_MAX_ITERATIONS", "5")) + + from factory.optimization.gate import evaluate_gate + + if not history: + print("HALT: no history") + return 2 + + latest = history[-1] + gate = evaluate_gate( + candidate_score=latest["score_end"], + current_score=latest["score_start"], + best_score=state["best_score"], + best_step=state["best_step"], + global_step=state["step"], + ) + + mutation_count = len([h for h in history if h["step"] > 1]) + if mutation_count >= max_iterations: + print(f"PROCEED: max iterations ({max_iterations}) reached") + return 0 + + if gate.accepted: + print(f"PROCEED: {gate.reason}") + return 0 + else: + print(f"RELOOP: {gate.reason}") + return 1 diff --git a/factory/optimization/loop.py b/factory/optimization/loop.py index d31059d5d..04a279ea6 100644 --- a/factory/optimization/loop.py +++ b/factory/optimization/loop.py @@ -49,6 +49,9 @@ def __init__( mutator: Mutator, config: LoopConfig | None = None, ) -> None: + log.warning( + "OptimizationLoop is deprecated; use 'factory optimize' (workflow mode) instead", + ) self.project_dir = Path(project_dir).resolve() self.surface = surface self.executor = executor diff --git a/factory/workflow/definitions.py b/factory/workflow/definitions.py index bc51efd7d..756595652 100644 --- a/factory/workflow/definitions.py +++ b/factory/workflow/definitions.py @@ -65,6 +65,7 @@ "evolve_workflow", "plan_workflow", "setup_eval_workflow", + "optimize_workflow", "register_all", "_get_builtin_registry", ] @@ -3765,6 +3766,97 @@ def trigger(state: ProjectState, ctx: dict[str, Any]) -> bool: ) +# ── W₁₇: Optimize Mode ──────────────────────────────────────────── + + +def optimize_workflow() -> Workflow: + """W₁₇: Optimize — inner-outer prompt optimization via workflow graph. + + baseline(FnNode) → gate_baseline(GateNode) → mutate(AgentNode) → + apply(FnNode) → execute(FnNode) → gate_improve(GateNode, RELOOP max 5) → + test_eval(FnNode) + """ + nodes: dict[str, Any] = {} + edges: list[Edge] = [] + + nodes["baseline"] = FnNode( + id="baseline", + command="factory optimize-step run-dev --project {project_path}", + writes={".factory/optimization/baseline.json"}, + ) + + nodes["gate_baseline"] = GateNode( + id="gate_baseline", + evaluator_type="fn", + evaluator_command="factory optimize-step check-gate --project {project_path} --baseline", + reads={".factory/optimization/baseline.json"}, + ) + + nodes["mutate"] = AgentNode( + id="mutate", + role=AgentRole.STRATEGIST, + prompt_template=( + "You are optimizing a skill prompt for a benchmark. " + "Read the current skill at .factory/optimization/current_skill.md. " + "Read the optimization state at .factory/optimization/state.json. " + "Read the latest step results from .factory/optimization/steps/ to identify failure patterns. " + "Analyze which tasks failed and why. Propose new rules to improve accuracy. " + "Output ONLY a JSON object: {\"rules\": [\"rule 1\", \"rule 2\", ...], \"reasoning\": \"...\"} " + "Write your output to .factory/optimization/mutation.json." + ), + reads={ + ".factory/optimization/current_skill.md", + ".factory/optimization/state.json", + }, + writes={".factory/optimization/mutation.json"}, + ) + + nodes["apply"] = FnNode( + id="apply", + command="factory optimize-step apply-patch --project {project_path}", + reads={".factory/optimization/mutation.json"}, + writes={".factory/optimization/current_skill.md"}, + ) + + nodes["execute"] = FnNode( + id="execute", + command="factory optimize-step run-dev --project {project_path}", + writes={".factory/optimization/state.json"}, + ) + + nodes["gate_improve"] = GateNode( + id="gate_improve", + evaluator_type="fn", + evaluator_command="factory optimize-step check-gate --project {project_path}", + reads={".factory/optimization/state.json"}, + ) + + nodes["test_eval"] = FnNode( + id="test_eval", + command="factory optimize-step run-test --project {project_path}", + writes={".factory/optimization/test_result.json"}, + ) + + edges = [ + Edge(source="baseline", target="gate_baseline"), + Edge(source="gate_baseline", target="mutate", condition=VerdictType.PROCEED), + Edge(source="gate_baseline", target="test_eval", condition=VerdictType.HALT), + Edge(source="mutate", target="apply"), + Edge(source="apply", target="execute"), + Edge(source="execute", target="gate_improve"), + Edge(source="gate_improve", target="test_eval", condition=VerdictType.PROCEED), + Edge(source="gate_improve", target="mutate", condition=VerdictType.RELOOP), + ] + + return Workflow( + name="optimize", + nodes=nodes, + edges=edges, + start_node="baseline", + terminal=True, + ) + + # ── Registry ───────────────────────────────────────────────────── _BUILTIN_REGISTRY: dict[str, Any] | None = None @@ -3798,6 +3890,7 @@ def _get_builtin_registry() -> dict[str, Any]: "plan": plan_workflow, "evolve": evolve_workflow, "setup-eval": setup_eval_workflow, + "optimize": optimize_workflow, "deep-qa": lambda: __import__( "factory.workflow.deep_qa", fromlist=["workflow"] ).workflow(), diff --git a/tests/test_cli_optimize.py b/tests/test_cli_optimize.py index c4ee3f5b4..f3183a206 100644 --- a/tests/test_cli_optimize.py +++ b/tests/test_cli_optimize.py @@ -95,11 +95,28 @@ def test_benchmark_auto_no_dir_errors(self, tmp_path) -> None: git_ref=None, docker_host=None, model="sonnet", + split_seed=42, + splits_dir=None, + legacy=True, ) result = cmd_optimize(args) assert result == 1 +class TestLegacyFlag: + """Verify --legacy flag routing.""" + + def test_legacy_flag_parsed(self) -> None: + parser = build_parser() + args = parser.parse_args(["optimize", "/tmp/proj", "--legacy"]) + assert args.legacy is True + + def test_legacy_flag_default_false(self) -> None: + parser = build_parser() + args = parser.parse_args(["optimize", "/tmp/proj"]) + assert args.legacy is False + + class TestCmdOptimize: """Test cmd_optimize with mocked dependencies.""" @@ -114,14 +131,16 @@ def test_missing_path_returns_1(self, tmp_path) -> None: git_ref=None, docker_host=None, model="sonnet", + legacy=False, ) result = cmd_optimize(args) assert result == 1 - def test_invalid_benchmark_returns_1(self, tmp_path) -> None: + def test_legacy_invalid_benchmark_returns_1(self, tmp_path) -> None: args = argparse.Namespace( path=str(tmp_path), benchmark="invalid_benchmark", + benchmark_dir=None, skill_path=None, steps=1, epochs=1, @@ -129,11 +148,14 @@ def test_invalid_benchmark_returns_1(self, tmp_path) -> None: git_ref=None, docker_host=None, model="sonnet", + split_seed=42, + splits_dir=None, + legacy=True, ) result = cmd_optimize(args) assert result == 1 - def test_successful_run_returns_0(self, tmp_path) -> None: + def test_legacy_successful_run_returns_0(self, tmp_path) -> None: from factory.optimization.loop import TrainResult from factory.optimization.types import StepRecord @@ -150,6 +172,7 @@ def test_successful_run_returns_0(self, tmp_path) -> None: args = argparse.Namespace( path=str(tmp_path), benchmark="searchqa", + benchmark_dir=None, skill_path=None, steps=2, epochs=1, @@ -157,9 +180,12 @@ def test_successful_run_returns_0(self, tmp_path) -> None: git_ref="main", docker_host=None, model="sonnet", + split_seed=42, + splits_dir=None, + legacy=True, ) - with patch("factory.optimization.loop.OptimizationLoop") as mock_loop_cls, \ + with patch("factory.optimization.OptimizationLoop") as mock_loop_cls, \ patch("factory.optimization.benchmarks.harbor.HarborBenchmark"), \ patch("factory.optimization.benchmarks.searchqa.SearchQAEvaluator"), \ patch("factory.optimization.mutators.agentic.AgenticMutator"): @@ -168,7 +194,33 @@ def test_successful_run_returns_0(self, tmp_path) -> None: assert result == 0 - def test_skill_path_loaded(self, tmp_path) -> None: + def test_workflow_mode_writes_state(self, tmp_path) -> None: + """Non-legacy mode writes initial state files and calls workflow.""" + args = argparse.Namespace( + path=str(tmp_path), + benchmark="searchqa", + skill_path=None, + steps=1, + concurrency=5, + git_ref="main", + docker_host=None, + model="sonnet", + legacy=False, + ) + + import json + + with patch("subprocess.run") as mock_run: + mock_run.return_value.returncode = 0 + # Pre-write state so summary reads succeed + opt_dir = tmp_path / ".factory" / "optimization" + result = cmd_optimize(args) + assert (opt_dir / "current_skill.md").exists() + state = json.loads((opt_dir / "state.json").read_text()) + assert state["step"] == 0 + assert state["history"] == [] + + def test_legacy_skill_path_loaded(self, tmp_path) -> None: from factory.optimization.loop import TrainResult skill_file = tmp_path / "skill.md" @@ -179,6 +231,7 @@ def test_skill_path_loaded(self, tmp_path) -> None: args = argparse.Namespace( path=str(tmp_path), benchmark="searchqa", + benchmark_dir=None, skill_path=str(skill_file), steps=1, epochs=1, @@ -186,9 +239,12 @@ def test_skill_path_loaded(self, tmp_path) -> None: git_ref="main", docker_host=None, model="sonnet", + split_seed=42, + splits_dir=None, + legacy=True, ) - with patch("factory.optimization.loop.OptimizationLoop") as mock_loop_cls, \ + with patch("factory.optimization.OptimizationLoop") as mock_loop_cls, \ patch("factory.optimization.benchmarks.harbor.HarborBenchmark"), \ patch("factory.optimization.benchmarks.searchqa.SearchQAEvaluator"), \ patch("factory.optimization.mutators.agentic.AgenticMutator"): diff --git a/tests/test_optimize_step.py b/tests/test_optimize_step.py new file mode 100644 index 000000000..45bbffb6e --- /dev/null +++ b/tests/test_optimize_step.py @@ -0,0 +1,237 @@ +"""Tests for factory.cli.optimize_step — workflow node helpers.""" + +from __future__ import annotations + +import argparse +import json + +import pytest + +from factory.cli.optimize_step import ( + cmd_optimize_step_apply_patch, + cmd_optimize_step_check_gate, + _read_state, + _write_state, +) + + +class TestApplyPatch: + """Test apply-patch JSON parsing and skill mutation.""" + + def test_valid_json(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + (opt_dir / "current_skill.md").write_text("# Skill\n\nBase content.\n") + (opt_dir / "mutation.json").write_text(json.dumps({ + "rules": ["Always check the question type", "Look for named entities"], + "reasoning": "test", + })) + + args = argparse.Namespace(project=str(tmp_path)) + result = cmd_optimize_step_apply_patch(args) + assert result == 0 + + skill = (opt_dir / "current_skill.md").read_text() + assert "Always check the question type" in skill + assert "Look for named entities" in skill + assert "## Learned Rules" in skill + + def test_markdown_wrapped_json(self, tmp_path) -> None: + """Strategist wraps JSON in markdown code blocks — regex fallback should handle it.""" + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + (opt_dir / "current_skill.md").write_text("# Skill\n") + (opt_dir / "mutation.json").write_text( + 'Here are the rules:\n```json\n{"rules": ["Rule A"], "reasoning": "test"}\n```\n' + ) + + args = argparse.Namespace(project=str(tmp_path)) + result = cmd_optimize_step_apply_patch(args) + assert result == 0 + + skill = (opt_dir / "current_skill.md").read_text() + assert "Rule A" in skill + + def test_empty_rules(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + (opt_dir / "current_skill.md").write_text("# Skill\n") + (opt_dir / "mutation.json").write_text(json.dumps({"rules": [], "reasoning": "nothing"})) + + args = argparse.Namespace(project=str(tmp_path)) + result = cmd_optimize_step_apply_patch(args) + assert result == 0 + + def test_missing_mutation_file(self, tmp_path) -> None: + args = argparse.Namespace(project=str(tmp_path)) + result = cmd_optimize_step_apply_patch(args) + assert result == 1 + + def test_unparseable_json(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + (opt_dir / "mutation.json").write_text("this is not json at all") + + args = argparse.Namespace(project=str(tmp_path)) + result = cmd_optimize_step_apply_patch(args) + assert result == 1 + + +class TestCheckGate: + """Test check-gate verdict logic.""" + + def test_baseline_positive_score(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + (opt_dir / "baseline.json").write_text(json.dumps({"score": 0.5})) + + args = argparse.Namespace(project=str(tmp_path), baseline=True) + result = cmd_optimize_step_check_gate(args) + assert result == 0 # PROCEED + + def test_baseline_zero_score(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + (opt_dir / "baseline.json").write_text(json.dumps({"score": 0.0})) + + args = argparse.Namespace(project=str(tmp_path), baseline=True) + result = cmd_optimize_step_check_gate(args) + assert result == 2 # HALT + + def test_improvement_proceed(self, tmp_path, monkeypatch) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + state = { + "step": 2, + "current_score": 0.7, + "best_score": 0.7, + "best_step": 2, + "history": [ + {"step": 1, "score_start": 0.0, "score_end": 0.5, "score_delta": 0.5, "verdict": "pending"}, + {"step": 2, "score_start": 0.5, "score_end": 0.7, "score_delta": 0.2, "verdict": "pending"}, + ], + } + _write_state(tmp_path, state) + monkeypatch.setenv("FACTORY_OPT_MAX_ITERATIONS", "5") + + args = argparse.Namespace(project=str(tmp_path), baseline=False) + result = cmd_optimize_step_check_gate(args) + assert result == 0 # PROCEED (improvement found) + + def test_no_improvement_reloop(self, tmp_path, monkeypatch) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + state = { + "step": 2, + "current_score": 0.5, + "best_score": 0.5, + "best_step": 1, + "history": [ + {"step": 1, "score_start": 0.0, "score_end": 0.5, "score_delta": 0.5, "verdict": "pending"}, + {"step": 2, "score_start": 0.5, "score_end": 0.4, "score_delta": -0.1, "verdict": "pending"}, + ], + } + _write_state(tmp_path, state) + monkeypatch.setenv("FACTORY_OPT_MAX_ITERATIONS", "5") + + args = argparse.Namespace(project=str(tmp_path), baseline=False) + result = cmd_optimize_step_check_gate(args) + assert result == 1 # RELOOP + + def test_max_iterations_proceed(self, tmp_path, monkeypatch) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + history = [ + {"step": i, "score_start": 0.5, "score_end": 0.5, "score_delta": 0.0, "verdict": "pending"} + for i in range(1, 7) + ] + state = { + "step": 6, + "current_score": 0.5, + "best_score": 0.5, + "best_step": 1, + "history": history, + } + _write_state(tmp_path, state) + monkeypatch.setenv("FACTORY_OPT_MAX_ITERATIONS", "5") + + args = argparse.Namespace(project=str(tmp_path), baseline=False) + result = cmd_optimize_step_check_gate(args) + assert result == 0 # PROCEED (max iterations) + + def test_empty_history_halt(self, tmp_path) -> None: + opt_dir = tmp_path / ".factory" / "optimization" + opt_dir.mkdir(parents=True) + + state = {"step": 0, "current_score": 0.0, "best_score": 0.0, "best_step": 0, "history": []} + _write_state(tmp_path, state) + + args = argparse.Namespace(project=str(tmp_path), baseline=False) + result = cmd_optimize_step_check_gate(args) + assert result == 2 # HALT + + +class TestStateReadWrite: + """Test state.json read/write helpers.""" + + def test_write_and_read(self, tmp_path) -> None: + state = {"step": 1, "current_score": 0.5, "best_score": 0.5, "best_step": 1, "history": []} + _write_state(tmp_path, state) + loaded = _read_state(tmp_path) + assert loaded == state + + def test_read_missing_returns_default(self, tmp_path) -> None: + state = _read_state(tmp_path) + assert state["step"] == 0 + assert state["history"] == [] + + def test_append_only_history(self, tmp_path) -> None: + state = _read_state(tmp_path) + state["history"].append({"step": 1, "score_start": 0.0, "score_end": 0.5}) + _write_state(tmp_path, state) + + state = _read_state(tmp_path) + state["history"].append({"step": 2, "score_start": 0.5, "score_end": 0.7}) + _write_state(tmp_path, state) + + final = _read_state(tmp_path) + assert len(final["history"]) == 2 + + +class TestOptimizeStepParser: + """Verify argparse setup for optimize-step subcommand.""" + + def test_optimize_step_parser_exists(self) -> None: + from factory.cli._main import build_parser + parser = build_parser() + args = parser.parse_args(["optimize-step", "run-dev", "--project", "/tmp/p"]) + assert args.command == "optimize-step" + assert args.optimize_step_command == "run-dev" + assert args.project == "/tmp/p" + + def test_apply_patch_subcommand(self) -> None: + from factory.cli._main import build_parser + parser = build_parser() + args = parser.parse_args(["optimize-step", "apply-patch", "--project", "/tmp/p"]) + assert args.optimize_step_command == "apply-patch" + + def test_check_gate_baseline_flag(self) -> None: + from factory.cli._main import build_parser + parser = build_parser() + args = parser.parse_args(["optimize-step", "check-gate", "--project", "/tmp/p", "--baseline"]) + assert args.optimize_step_command == "check-gate" + assert args.baseline is True + + def test_run_test_subcommand(self) -> None: + from factory.cli._main import build_parser + parser = build_parser() + args = parser.parse_args(["optimize-step", "run-test", "--project", "/tmp/p"]) + assert args.optimize_step_command == "run-test" diff --git a/tests/test_workflow_optimize.py b/tests/test_workflow_optimize.py new file mode 100644 index 000000000..6089a4a5a --- /dev/null +++ b/tests/test_workflow_optimize.py @@ -0,0 +1,214 @@ +"""Tests for the optimize workflow graph definition.""" + +from __future__ import annotations + +from collections import defaultdict + +import pytest + +from factory.workflow.definitions import optimize_workflow +from factory.workflow.primitives import ( + AgentNode, + AgentRole, + FnNode, + GateNode, + VerdictType, +) + + +class TestOptimizeWorkflowStructure: + """Verify the optimize workflow graph has the correct structure.""" + + def test_has_7_nodes(self) -> None: + wf = optimize_workflow() + assert len(wf.nodes) == 7 + + def test_node_ids(self) -> None: + wf = optimize_workflow() + expected = {"baseline", "gate_baseline", "mutate", "apply", "execute", "gate_improve", "test_eval"} + assert set(wf.nodes.keys()) == expected + + def test_node_types(self) -> None: + wf = optimize_workflow() + assert isinstance(wf.nodes["baseline"], FnNode) + assert isinstance(wf.nodes["gate_baseline"], GateNode) + assert isinstance(wf.nodes["mutate"], AgentNode) + assert isinstance(wf.nodes["apply"], FnNode) + assert isinstance(wf.nodes["execute"], FnNode) + assert isinstance(wf.nodes["gate_improve"], GateNode) + assert isinstance(wf.nodes["test_eval"], FnNode) + + def test_mutate_is_strategist(self) -> None: + wf = optimize_workflow() + assert wf.nodes["mutate"].role == AgentRole.STRATEGIST + + def test_start_node(self) -> None: + wf = optimize_workflow() + assert wf.start_node == "baseline" + + def test_is_terminal(self) -> None: + wf = optimize_workflow() + assert wf.terminal is True + + def test_name(self) -> None: + wf = optimize_workflow() + assert wf.name == "optimize" + + +class TestOptimizeWorkflowEdges: + """Verify edge connectivity.""" + + def test_edge_count(self) -> None: + wf = optimize_workflow() + assert len(wf.edges) == 8 + + def test_baseline_to_gate_baseline(self) -> None: + wf = optimize_workflow() + edge = next(e for e in wf.edges if e.source == "baseline") + assert edge.target == "gate_baseline" + assert edge.condition is None + + def test_gate_baseline_proceed_to_mutate(self) -> None: + wf = optimize_workflow() + edge = next( + e for e in wf.edges + if e.source == "gate_baseline" and e.condition == VerdictType.PROCEED + ) + assert edge.target == "mutate" + + def test_gate_baseline_halt_to_test_eval(self) -> None: + wf = optimize_workflow() + edge = next( + e for e in wf.edges + if e.source == "gate_baseline" and e.condition == VerdictType.HALT + ) + assert edge.target == "test_eval" + + def test_mutate_to_apply(self) -> None: + wf = optimize_workflow() + edge = next(e for e in wf.edges if e.source == "mutate") + assert edge.target == "apply" + + def test_apply_to_execute(self) -> None: + wf = optimize_workflow() + edge = next(e for e in wf.edges if e.source == "apply") + assert edge.target == "execute" + + def test_execute_to_gate_improve(self) -> None: + wf = optimize_workflow() + edge = next(e for e in wf.edges if e.source == "execute") + assert edge.target == "gate_improve" + + def test_gate_improve_proceed_to_test_eval(self) -> None: + wf = optimize_workflow() + edge = next( + e for e in wf.edges + if e.source == "gate_improve" and e.condition == VerdictType.PROCEED + ) + assert edge.target == "test_eval" + + def test_gate_improve_reloop_to_mutate(self) -> None: + wf = optimize_workflow() + edge = next( + e for e in wf.edges + if e.source == "gate_improve" and e.condition == VerdictType.RELOOP + ) + assert edge.target == "mutate" + + +class TestOptimizeWorkflowValidation: + """Graph validation should pass.""" + + def test_validates_clean(self) -> None: + wf = optimize_workflow() + issues = wf.validate_graph() + assert issues == [], f"optimize workflow has issues: {issues}" + + +class TestOptimizeWorkflowRegistry: + """Verify optimize is in the builtin registry.""" + + def test_in_registry(self) -> None: + from factory.workflow.definitions import _get_builtin_registry + registry = _get_builtin_registry() + assert "optimize" in registry + + def test_registry_callable_returns_workflow(self) -> None: + from factory.workflow.definitions import _get_builtin_registry + registry = _get_builtin_registry() + wf = registry["optimize"]() + assert wf.name == "optimize" + assert len(wf.nodes) == 7 + + +class TestOptimizeWorkflowReadsWrites: + """Verify node reads/writes are consistent for file-based inter-node communication.""" + + def test_baseline_writes_baseline_json(self) -> None: + wf = optimize_workflow() + assert ".factory/optimization/baseline.json" in wf.nodes["baseline"].writes + + def test_mutate_reads_skill_and_state(self) -> None: + wf = optimize_workflow() + reads = wf.nodes["mutate"].reads + assert ".factory/optimization/current_skill.md" in reads + assert ".factory/optimization/state.json" in reads + + def test_mutate_writes_mutation(self) -> None: + wf = optimize_workflow() + assert ".factory/optimization/mutation.json" in wf.nodes["mutate"].writes + + def test_apply_reads_mutation(self) -> None: + wf = optimize_workflow() + assert ".factory/optimization/mutation.json" in wf.nodes["apply"].reads + + def test_apply_writes_skill(self) -> None: + wf = optimize_workflow() + assert ".factory/optimization/current_skill.md" in wf.nodes["apply"].writes + + def test_test_eval_writes_result(self) -> None: + wf = optimize_workflow() + assert ".factory/optimization/test_result.json" in wf.nodes["test_eval"].writes + + +class TestOptimizeWorkflowPaths: + """Verify the two main paths through the graph.""" + + def _build_adj(self, wf): + adj = defaultdict(list) + for e in wf.edges: + adj[e.source].append((e.target, e.condition)) + return adj + + def test_baseline_halt_path(self) -> None: + """baseline → gate_baseline → (HALT) → test_eval""" + wf = optimize_workflow() + adj = self._build_adj(wf) + # baseline → gate_baseline + targets = [t for t, c in adj["baseline"]] + assert "gate_baseline" in targets + # gate_baseline → test_eval on HALT + halt_targets = [t for t, c in adj["gate_baseline"] if c == VerdictType.HALT] + assert "test_eval" in halt_targets + + def test_mutation_reloop_proceed_path(self) -> None: + """gate_baseline → (PROCEED) → mutate → apply → execute → gate_improve → (PROCEED) → test_eval""" + wf = optimize_workflow() + adj = self._build_adj(wf) + proceed_from_baseline = [t for t, c in adj["gate_baseline"] if c == VerdictType.PROCEED] + assert "mutate" in proceed_from_baseline + + targets_from_mutate = [t for t, c in adj["mutate"]] + assert "apply" in targets_from_mutate + + targets_from_apply = [t for t, c in adj["apply"]] + assert "execute" in targets_from_apply + + targets_from_execute = [t for t, c in adj["execute"]] + assert "gate_improve" in targets_from_execute + + reloop_targets = [t for t, c in adj["gate_improve"] if c == VerdictType.RELOOP] + assert "mutate" in reloop_targets + + proceed_targets = [t for t, c in adj["gate_improve"] if c == VerdictType.PROCEED] + assert "test_eval" in proceed_targets