weco 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

weco/cli.py

@@ -72,6 +72,11 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Save execution output to .runs/<run-id>/outputs/step_<n>.out.txt with JSONL index",
     )
+    run_parser.add_argument(
+        "--apply-change",
+        action="store_true",
+        help="Automatically apply the best solution to the source file without prompting",
+    )
 
 
 def configure_credits_parser(credits_parser: argparse.ArgumentParser) -> None:
@@ -118,6 +123,11 @@ def configure_resume_parser(resume_parser: argparse.ArgumentParser) -> None:
     resume_parser.add_argument(
         "run_id", type=str, help="The UUID of the run to resume (e.g., '0002e071-1b67-411f-a514-36947f0c4b31')"
     )
+    resume_parser.add_argument(
+        "--apply-change",
+        action="store_true",
+        help="Automatically apply the best solution to the source file without prompting",
+    )
 
 
 def execute_run_command(args: argparse.Namespace) -> None:
@@ -136,6 +146,7 @@ def execute_run_command(args: argparse.Namespace) -> None:
         console=console,
         eval_timeout=args.eval_timeout,
         save_logs=args.save_logs,
+        apply_change=args.apply_change,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)
@@ -145,7 +156,7 @@ def execute_resume_command(args: argparse.Namespace) -> None:
     """Execute the 'weco resume' command with all its logic."""
     from .optimizer import resume_optimization
 
-    success = resume_optimization(run_id=args.run_id, console=console)
+    success = resume_optimization(run_id=args.run_id, console=console, apply_change=args.apply_change)
     sys.exit(0 if success else 1)
 
 

weco/optimizer.py

@@ -30,7 +30,7 @@ from .panels import (
     create_optimization_layout,
     create_end_optimization_layout,
 )
-from .utils import read_additional_instructions, read_from_path, write_to_path, run_evaluation, smooth_update
+from .utils import read_additional_instructions, read_from_path, write_to_path, run_evaluation_with_file_swap, smooth_update
 
 
 def save_execution_output(runs_dir: pathlib.Path, step: int, output: str) -> None:
@@ -134,6 +134,7 @@ def execute_optimization(
     console: Optional[Console] = None,
     eval_timeout: Optional[int] = None,
     save_logs: bool = False,
+    apply_change: bool = False,
 ) -> bool:
     """
     Execute the core optimization logic.
@@ -182,6 +183,8 @@ def execute_optimization(
     optimization_completed_normally = False
     user_stop_requested_flag = False
 
+    best_solution_code = None
+    original_source_code = None  # Make available to the finally block
     try:
         # --- Login/Authentication Handling (now mandatory) ---
         weco_api_key, auth_headers = handle_authentication(console)
@@ -209,6 +212,7 @@ def execute_optimization(
         processed_additional_instructions = read_additional_instructions(additional_instructions=additional_instructions)
         source_fp = pathlib.Path(source)
         source_code = read_from_path(fp=source_fp, is_json=False)
+        original_source_code = source_code
 
         # --- Panel Initialization ---
         summary_panel = SummaryPanel(maximize=maximize, metric_name=metric, total_steps=steps, model=model, runs_dir=log_dir)
@@ -272,10 +276,6 @@ def execute_optimization(
                 }
                 with open(jsonl_file, "w", encoding="utf-8") as f:
                     f.write(json.dumps(metadata) + "\n")
-            # Write the initial code string to the logs
-            write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=run_response["code"])
-            # Write the initial code string to the source file path
-            write_to_path(fp=source_fp, content=run_response["code"])
 
             # Update the panels with the initial solution
             # Add run id and run name now that we have it
@@ -307,6 +307,7 @@ def execute_optimization(
                 best_node=None,
             )
             current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=0)
+
             # Update the live layout with the initial solution panels
             smooth_update(
                 live=live,
@@ -321,8 +322,17 @@ def execute_optimization(
                 transition_delay=0.1,
             )
 
-            # Run evaluation on the initial solution
-            term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
+            # Write the initial code string to the logs
+            write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=run_response["code"])
+            # Run evaluation on the initial solution (file swap ensures original is restored)
+            term_out = run_evaluation_with_file_swap(
+                file_path=source_fp,
+                new_content=run_response["code"],
+                original_content=source_code,
+                eval_command=eval_command,
+                timeout=eval_timeout,
+            )
+
             # Save logs if requested
             if save_logs:
                 save_execution_output(runs_dir, step=0, output=term_out)
@@ -358,8 +368,7 @@ def execute_optimization(
                 )
                 # Save next solution (.runs/<run-id>/step_<step>.<extension>)
                 write_to_path(fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"])
-                # Write the next solution to the source file
-                write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
+
                 status_response = get_optimization_run_status(
                     console=console, run_id=run_id, include_history=True, auth_headers=auth_headers
                 )
@@ -395,7 +404,16 @@ def execute_optimization(
                     ],
                     transition_delay=0.08,  # Slightly longer delay for more noticeable transitions
                 )
-                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
+
+                # Run evaluation and restore original code after
+                term_out = run_evaluation_with_file_swap(
+                    file_path=source_fp,
+                    new_content=eval_and_next_solution_response["code"],
+                    original_content=source_code,
+                    eval_command=eval_command,
+                    timeout=eval_timeout,
+                )
+
                 # Save logs if requested
                 if save_logs:
                     save_execution_output(runs_dir, step=step, output=term_out)
@@ -444,15 +462,14 @@ def execute_optimization(
                 # Save optimization results
                 # If the best solution does not exist or is has not been measured at the end of the optimization
                 # save the original solution as the best solution
-                if best_solution_node is not None:
-                    best_solution_content = best_solution_node.code
-                else:
-                    best_solution_content = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
+                try:
+                    best_solution_code = best_solution_node.code
+                except AttributeError:
+                    best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
 
                 # Save best solution to .runs/<run-id>/best.<extension>
-                write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
-                # write the best solution to the source file
-                write_to_path(fp=source_fp, content=best_solution_content)
+                write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
+
                 # Mark as completed normally for the finally block
                 optimization_completed_normally = True
                 live.update(end_optimization_layout)
@@ -491,6 +508,21 @@ def execute_optimization(
                     else "CLI terminated unexpectedly without a specific exception captured."
                 )
 
+            # raise Exception(best_solution_code, original_source_code)
+            if best_solution_code and best_solution_code != original_source_code:
+                # Determine whether to apply: automatically if --apply-change is set, otherwise ask user
+                should_apply = apply_change or summary_panel.ask_user_feedback(
+                    live=live,
+                    layout=end_optimization_layout,
+                    question="Would you like to apply the best solution to the source file?",
+                    default=True,
+                )
+                if should_apply:
+                    write_to_path(fp=source_fp, content=best_solution_code)
+                    console.print("[green]Best solution applied to the source file.[/]\n")
+            else:
+                console.print("[green]A better solution was not found. No changes to apply.[/]\n")
+
             report_termination(
                 run_id=run_id,
                 status_update=status,
@@ -507,7 +539,7 @@ def execute_optimization(
     return optimization_completed_normally or user_stop_requested_flag
 
 
-def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
+def resume_optimization(run_id: str, console: Optional[Console] = None, apply_change: bool = False) -> bool:
     """Resume an interrupted run from the most recent node and continue optimization."""
     if console is None:
         console = Console()
@@ -543,6 +575,9 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
     optimization_completed_normally = False
     user_stop_requested_flag = False
 
+    best_solution_code = None
+    original_source_code = None
+
     try:
         # --- Login/Authentication Handling (now mandatory) ---
         weco_api_key, auth_headers = handle_authentication(console)
@@ -619,9 +654,12 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
         save_logs = bool(resume_resp.get("save_logs", False))
         eval_timeout = resume_resp.get("eval_timeout")
 
-        # Write last solution code to source path
+        # Read the original source code from the file before we start modifying it
         source_fp = pathlib.Path(source_path)
         source_fp.parent.mkdir(parents=True, exist_ok=True)
+        # Store the original content to restore after each evaluation
+        original_source_code = read_from_path(fp=source_fp, is_json=False) if source_fp.exists() else ""
+
         code_to_restore = resume_resp.get("code") or resume_resp.get("source_code") or ""
         write_to_path(fp=source_fp, content=code_to_restore)
 
@@ -689,7 +727,13 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
 
             # If missing output, evaluate once before first suggest
             if term_out is None or len(term_out.strip()) == 0:
-                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
+                term_out = run_evaluation_with_file_swap(
+                    file_path=source_fp,
+                    new_content=code_to_restore,
+                    original_content=original_source_code,
+                    eval_command=eval_command,
+                    timeout=eval_timeout,
+                )
                 eval_output_panel.update(output=term_out)
                 # Update the evaluation output panel
                 smooth_update(
@@ -727,9 +771,8 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                     auth_headers=auth_headers,
                 )
 
-                # Save next solution file(s)
+                # Save next solution to logs
                 write_to_path(fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"])
-                write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
 
                 # Refresh status with history and update panels
                 status_response = get_optimization_run_status(
@@ -744,6 +787,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                 current_solution_node = get_node_from_status(
                     status_response=status_response, solution_id=eval_and_next_solution_response["solution_id"]
                 )
+
                 solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
                 current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
                 eval_output_panel.clear()
@@ -760,8 +804,14 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                     transition_delay=0.08,
                 )
 
-                # Evaluate this new solution
-                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
+                # Evaluate this new solution and restore original code after
+                term_out = run_evaluation_with_file_swap(
+                    file_path=source_fp,
+                    new_content=eval_and_next_solution_response["code"],
+                    original_content=original_source_code,
+                    eval_command=eval_command,
+                    timeout=eval_timeout,
+                )
                 if save_logs:
                     save_execution_output(runs_dir, step=step, output=term_out)
                 eval_output_panel.update(output=term_out)
@@ -801,13 +851,13 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                 end_optimization_layout["best_solution"].update(best_solution_panel)
 
                 # Save best
-                if best_solution_node is not None:
-                    best_solution_content = best_solution_node.code
-                else:
-                    best_solution_content = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
+                try:
+                    best_solution_code = best_solution_node.code
+                except AttributeError:
+                    best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
+
+                write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
 
-                write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
-                write_to_path(fp=source_fp, content=best_solution_content)
                 optimization_completed_normally = True
                 live.update(end_optimization_layout)
 
@@ -826,7 +876,11 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
         if heartbeat_thread and heartbeat_thread.is_alive():
             heartbeat_thread.join(timeout=2)
 
-        run_id = resume_resp.get("run_id")
+        try:
+            run_id = resume_resp.get("run_id")
+        except Exception:
+            run_id = None
+
         # Report final status if run exists
         if run_id:
             if optimization_completed_normally:
@@ -840,6 +894,20 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                     if "e" in locals() and isinstance(locals()["e"], Exception)
                     else "CLI terminated unexpectedly without a specific exception captured."
                 )
+
+            if best_solution_code and best_solution_code != original_source_code:
+                should_apply = apply_change or summary_panel.ask_user_feedback(
+                    live=live,
+                    layout=end_optimization_layout,
+                    question="Would you like to apply the best solution to the source file?",
+                    default=True,
+                )
+                if should_apply:
+                    write_to_path(fp=source_fp, content=best_solution_code)
+                    console.print("[green]Best solution applied to the source file.[/]\n")
+            else:
+                console.print("[green]A better solution was not found. No changes to apply.[/]\n")
+
             report_termination(
                 run_id=run_id,
                 status_update=status,

weco/panels.py

@@ -5,6 +5,10 @@ from rich.layout import Layout
 from rich.panel import Panel
 from rich.syntax import Syntax
 from rich import box
+from rich.console import Console
+from rich.live import Live
+
+from rich.prompt import Confirm
 from typing import Dict, List, Optional, Union, Tuple
 from pathlib import Path
 from .__init__ import __dashboard_url__
@@ -22,6 +26,7 @@ class SummaryPanel:
         runs_dir: str,
         run_id: str = None,
         run_name: str = None,
+        console: Optional[Console] = None,
     ):
         self.maximize = maximize
         self.metric_name = metric_name
@@ -32,6 +37,8 @@ class SummaryPanel:
         self.run_name = run_name if run_name is not None else "N/A"
         self.dashboard_url = "N/A"
         self.thinking_content = ""
+        self.user_input = ""
+        self.console = Console()
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=20),
@@ -67,6 +74,25 @@ class SummaryPanel:
         """Clear the thinking content."""
         self.thinking_content = ""
 
+    def ask_user_feedback(self, live: Live, layout: Layout, question: str, default: bool = True) -> bool:
+        """
+        Ask a yes/no question while keeping the main layout fixed.
+        Uses Rich's Confirm for a clean user experience.
+        """
+        # Stop live updates temporarily to prevent layout from moving
+        live.stop()
+
+        try:
+            # Use Rich's built-in Confirm
+            result = Confirm.ask(question, default=default)
+        except (KeyboardInterrupt, EOFError):
+            result = default
+        finally:
+            # Resume live updates
+            live.start()
+
+        return result
+
     def get_display(self, final_message: Optional[str] = None) -> Panel:
         """Return a Rich panel summarising the current run."""
         # ───────────────────── summary grid ──────────────────────

weco/utils.py

@@ -106,6 +106,41 @@ def truncate_output(output: str) -> str:
         return output
 
 
+def run_evaluation_with_file_swap(
+    file_path: pathlib.Path, new_content: str, original_content: str, eval_command: str, timeout: int | None = None
+) -> str:
+    """
+    Temporarily write new content to a file, run evaluation, then restore original.
+
+    This function ensures the file is always restored to its original state,
+    even if an exception occurs during evaluation.
+
+    Args:
+        file_path: Path to the file to temporarily modify
+        new_content: The new content to write for evaluation
+        original_content: The original content to restore after evaluation
+        eval_command: The shell command to run for evaluation
+        timeout: Optional timeout for the evaluation command
+
+    Returns:
+        The output from running the evaluation command
+
+    Raises:
+        Any exception raised by run_evaluation will be re-raised after
+        the file is restored to its original state.
+    """
+    # Write the new content
+    write_to_path(fp=file_path, content=new_content)
+
+    try:
+        # Run the evaluation
+        output = run_evaluation(eval_command=eval_command, timeout=timeout)
+        return output
+    finally:
+        # Always restore the original file, even if evaluation fails
+        write_to_path(fp=file_path, content=original_content)
+
+
 def run_evaluation(eval_command: str, timeout: int | None = None) -> str:
     """Run the evaluation command on the code and return the output."""
     process = subprocess.Popen(

{weco-0.3.3.dist-info → weco-0.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.3.3
+Version: 0.3.4
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: 
@@ -322,6 +322,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
 | `-l, --log-dir`                | Path to the directory to log intermediate steps and final optimization result.                                                                                                                                             | `.runs/`                                                                                                                                               | `-l ./logs/`        |
 | `--eval-timeout`       | Timeout in seconds for each step in evaluation.                                                                                                                                                                             | No timeout (unlimited)                                                                                                                                                  | `--eval-timeout 3600`             |
 | `--save-logs`          | Save execution output from each optimization step to disk. Creates timestamped directories with raw output files and a JSONL index for tracking execution history.                                                        | `False`                                                                                                                                                 | `--save-logs`       |
+| `--apply-change`       | Automatically apply the best solution to the source file without prompting.                                                                                                                                                | `False`                                                                                                                                                 | `--apply-change`       |
 
 ---
 
@@ -375,6 +376,7 @@ Arguments for `weco resume`:
 | Argument | Description | Example |
 |----------|-------------|---------|
 | `run-id` | The UUID of the run to resume (shown at the start of each run) | `0002e071-1b67-411f-a514-36947f0c4b31` |
+| `--apply-change` | Automatically apply the best solution to the source file without prompting | `--apply-change` |
 
 Notes:
 - Works only for interrupted runs (status: `error`, `terminated`, etc.).

weco-0.3.4.dist-info/RECORD

@@ -0,0 +1,16 @@
+weco/__init__.py,sha256=ClO0uT6GKOA0iSptvP0xbtdycf0VpoPTq37jHtvlhtw,303
+weco/api.py,sha256=zKcI4riwruK6CjV_vcL8RlsJGRXO40iP0WxeETtzPIY,18430
+weco/auth.py,sha256=O31Hoj-Loi8DWJJG2LfeWgUMuNqAUeGDpd2ZGjA9Ah0,9997
+weco/chatbot.py,sha256=EIK2WaOul9gn_yHLThjsZV7RnE8t3XQPwgRkO5tybSU,38415
+weco/cli.py,sha256=7Q6ZwkD8N2x_PRovaFkF9OvtKkKX3hs_CobvS8xKXkQ,11535
+weco/constants.py,sha256=V6yFugTznKm5EC2_jr4I_whd7sqI80HiPggRn0az580,406
+weco/credits.py,sha256=C08x-TRcLg3ccfKqMGNRY7zBn7t3r7LZ119bxgfztaI,7629
+weco/optimizer.py,sha256=fvxpKHeET84vIBWEM_50iCInbxMYs0e44sJqe7zVzK0,43848
+weco/panels.py,sha256=OLss2PLgJqdo9w-Gq0TozyE42lsJvgjtqmfyCHpHGtA,16966
+weco/utils.py,sha256=SMFh7ngo6gsjb2g20wpPe0O0T8xR7U_7X4-kqxJHkGc,8308
+weco-0.3.4.dist-info/licenses/LICENSE,sha256=9LUfoGHjLPtak2zps2kL2tm65HAZIICx_FbLaRuS4KU,11337
+weco-0.3.4.dist-info/METADATA,sha256=TGIKIIAl_0xQ9UBIjr3opE45ni_ZImzhLe29aBWQo3Q,29244
+weco-0.3.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+weco-0.3.4.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.3.4.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.3.4.dist-info/RECORD,,

weco-0.3.3.dist-info/RECORD

@@ -1,16 +0,0 @@
-weco/__init__.py,sha256=ClO0uT6GKOA0iSptvP0xbtdycf0VpoPTq37jHtvlhtw,303
-weco/api.py,sha256=zKcI4riwruK6CjV_vcL8RlsJGRXO40iP0WxeETtzPIY,18430
-weco/auth.py,sha256=O31Hoj-Loi8DWJJG2LfeWgUMuNqAUeGDpd2ZGjA9Ah0,9997
-weco/chatbot.py,sha256=EIK2WaOul9gn_yHLThjsZV7RnE8t3XQPwgRkO5tybSU,38415
-weco/cli.py,sha256=5TusCKQ8o4CUdqxOSkvWyjMzt86O9sj21SJoeU0Ni5w,11098
-weco/constants.py,sha256=V6yFugTznKm5EC2_jr4I_whd7sqI80HiPggRn0az580,406
-weco/credits.py,sha256=C08x-TRcLg3ccfKqMGNRY7zBn7t3r7LZ119bxgfztaI,7629
-weco/optimizer.py,sha256=nOKFmwPdFLcQ7RF4ielsD7iRfPMxvGr07pS9ocbW9C8,41282
-weco/panels.py,sha256=fnGPtmvxpx21AuBCtCFu1f_BpSxybNr2lhjIIKIutrY,16133
-weco/utils.py,sha256=erDDrA_g3KSlel6YEAGALlV_k8ftT-VQnPT1BrmzK8k,7021
-weco-0.3.3.dist-info/licenses/LICENSE,sha256=9LUfoGHjLPtak2zps2kL2tm65HAZIICx_FbLaRuS4KU,11337
-weco-0.3.3.dist-info/METADATA,sha256=aAE9oMp_SKJiQSG189k7fsb5ovvidj2t1DTuqNfvc68,28700
-weco-0.3.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-weco-0.3.3.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.3.3.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.3.3.dist-info/RECORD,,