weco 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

weco/optimizer.py

@@ -20,7 +20,6 @@ from .api import (
 from .auth import handle_authentication
 from .panels import (
     SummaryPanel,
-    PlanPanel,
     Node,
     MetricTreePanel,
     EvaluationOutputPanel,
@@ -37,6 +36,7 @@ from .utils import (
     smooth_update,
     format_number,
 )
+from .constants import DEFAULT_API_TIMEOUT
 
 
 # --- Heartbeat Sender Class ---
@@ -63,7 +63,7 @@ class HeartbeatSender(threading.Thread):
 
         except Exception as e:
             # Catch any unexpected error in the loop to prevent silent thread death
-            print(f"[ERROR HeartbeatSender] Unhandled exception in run loop for run {self.run_id}: {e}", file=sys.stderr)
+            print(f"[ERROR HeartbeatSender] Unexpected error in heartbeat thread for run {self.run_id}: {e}", file=sys.stderr)
             traceback.print_exc(file=sys.stderr)
             # The loop will break due to the exception, and thread will terminate via finally.
 
@@ -78,6 +78,7 @@ def execute_optimization(
     log_dir: str = ".runs",
     additional_instructions: Optional[str] = None,
     console: Optional[Console] = None,
+    eval_timeout: Optional[int] = None,
 ) -> bool:
     """
     Execute the core optimization logic.
@@ -153,14 +154,13 @@ def execute_optimization(
             "debug_prob": 0.5,
             "max_debug_depth": max(1, math.ceil(0.1 * steps)),
         }
-        timeout = 800
+        api_timeout = DEFAULT_API_TIMEOUT
         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)
 
         # --- Panel Initialization ---
         summary_panel = SummaryPanel(maximize=maximize, metric_name=metric, total_steps=steps, model=model, runs_dir=log_dir)
-        plan_panel = PlanPanel()
         solution_panels = SolutionPanels(metric_name=metric, source_fp=source_fp)
         eval_output_panel = EvaluationOutputPanel()
         tree_panel = MetricTreePanel(maximize=maximize)
@@ -181,9 +181,10 @@ def execute_optimization(
             additional_instructions=processed_additional_instructions,
             api_keys=llm_api_keys,
             auth_headers=auth_headers,
-            timeout=timeout,
+            timeout=api_timeout,
         )
         run_id = run_response["run_id"]
+        run_name = run_response["run_name"]
         current_run_id_for_heartbeat = run_id
 
         # --- Start Heartbeat Thread ---
@@ -203,12 +204,14 @@ def execute_optimization(
             write_to_path(fp=source_fp, content=run_response["code"])
 
             # Update the panels with the initial solution
-            summary_panel.set_run_id(run_id=run_id)  # Add run id now that we have it
+            # Add run id and run name now that we have it
+            summary_panel.set_run_id(run_id=run_id)
+            summary_panel.set_run_name(run_name=run_name)
             # Set the step of the progress bar
             summary_panel.set_step(step=0)
             # Update the token counts
             summary_panel.update_token_counts(usage=run_response["usage"])
-            plan_panel.update(plan=run_response["plan"])
+            summary_panel.update_thinking(thinking=run_response["plan"])
             # Build the metric tree
             tree_panel.build_metric_tree(
                 nodes=[
@@ -218,7 +221,7 @@ def execute_optimization(
                         "code": run_response["code"],
                         "step": 0,
                         "metric_value": None,
-                        "is_buggy": False,
+                        "is_buggy": None,
                     }
                 ]
             )
@@ -227,7 +230,7 @@ def execute_optimization(
             # Update the solution panels with the initial solution and get the panel displays
             solution_panels.update(
                 current_node=Node(
-                    id=run_response["solution_id"], parent_id=None, code=run_response["code"], metric=None, is_buggy=False
+                    id=run_response["solution_id"], parent_id=None, code=run_response["code"], metric=None, is_buggy=None
                 ),
                 best_node=None,
             )
@@ -238,7 +241,6 @@ def execute_optimization(
                 layout=layout,
                 sections_to_update=[
                     ("summary", summary_panel.get_display()),
-                    ("plan", plan_panel.get_display()),
                     ("tree", tree_panel.get_display(is_done=False)),
                     ("current_solution", current_solution_panel),
                     ("best_solution", best_solution_panel),
@@ -248,7 +250,7 @@ def execute_optimization(
             )
 
             # Run evaluation on the initial solution
-            term_out = run_evaluation(eval_command=eval_command)
+            term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
             # Update the evaluation output panel
             eval_output_panel.update(output=term_out)
             smooth_update(
@@ -265,7 +267,7 @@ def execute_optimization(
                 if run_id:
                     try:
                         current_status_response = get_optimization_run_status(
-                            run_id=run_id, include_history=False, timeout=30, auth_headers=auth_headers
+                            console=console, run_id=run_id, include_history=False, timeout=(10, 30), auth_headers=auth_headers
                         )
                         current_run_status_val = current_status_response.get("status")
                         if current_run_status_val == "stopping":
@@ -273,30 +275,31 @@ def execute_optimization(
                             user_stop_requested_flag = True
                             break
                     except requests.exceptions.RequestException as e:
-                        console.print(f"\n[bold red]Warning: Could not check run status: {e}. Continuing optimization...[/]")
+                        console.print(f"\n[bold red]Warning: Unable to check run status: {e}. Continuing optimization...[/]")
                     except Exception as e:
                         console.print(f"\n[bold red]Warning: Error checking run status: {e}. Continuing optimization...[/]")
 
                 # Send feedback and get next suggestion
                 eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
+                    console=console,
                     run_id=run_id,
                     execution_output=term_out,
                     additional_instructions=current_additional_instructions,
                     api_keys=llm_api_keys,
                     auth_headers=auth_headers,
-                    timeout=timeout,
+                    timeout=api_timeout,
                 )
                 # 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(
-                    run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    console=console, run_id=run_id, include_history=True, timeout=api_timeout, auth_headers=auth_headers
                 )
                 # Update the step of the progress bar, token counts, plan and metric tree
                 summary_panel.set_step(step=step)
                 summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
-                plan_panel.update(plan=eval_and_next_solution_response["plan"])
+                summary_panel.update_thinking(thinking=eval_and_next_solution_response["plan"])
 
                 nodes_list_from_status = status_response.get("nodes")
                 tree_panel.build_metric_tree(nodes=nodes_list_from_status if nodes_list_from_status is not None else [])
@@ -327,7 +330,9 @@ def execute_optimization(
                                 is_buggy=node_data["is_buggy"],
                             )
                 if current_solution_node is None:
-                    raise ValueError("Current solution node not found in nodes list from status response")
+                    raise ValueError(
+                        "Current solution node not found in the optimization status response. This may indicate a synchronization issue with the backend."
+                    )
 
                 # Update the solution panels with the current and best solution
                 solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
@@ -339,7 +344,6 @@ def execute_optimization(
                     layout=layout,
                     sections_to_update=[
                         ("summary", summary_panel.get_display()),
-                        ("plan", plan_panel.get_display()),
                         ("tree", tree_panel.get_display(is_done=False)),
                         ("current_solution", current_solution_panel),
                         ("best_solution", best_solution_panel),
@@ -347,7 +351,7 @@ def execute_optimization(
                     ],
                     transition_delay=0.08,  # Slightly longer delay for more noticeable transitions
                 )
-                term_out = run_evaluation(eval_command=eval_command)
+                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
                 eval_output_panel.update(output=term_out)
                 smooth_update(
                     live=live,
@@ -361,17 +365,18 @@ def execute_optimization(
                 current_additional_instructions = read_additional_instructions(additional_instructions=additional_instructions)
                 # Evaluate the final solution thats been generated
                 eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
+                    console=console,
                     run_id=run_id,
                     execution_output=term_out,
                     additional_instructions=current_additional_instructions,
                     api_keys=llm_api_keys,
-                    timeout=timeout,
+                    timeout=api_timeout,
                     auth_headers=auth_headers,
                 )
                 summary_panel.set_step(step=steps)
                 summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
                 status_response = get_optimization_run_status(
-                    run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    console=console, run_id=run_id, include_history=True, timeout=api_timeout, auth_headers=auth_headers
                 )
                 # No need to update the plan panel since we have finished the optimization
                 # Get the optimization run status for

weco/panels.py

@@ -7,30 +7,40 @@ from rich.syntax import Syntax
 from rich import box
 from typing import Dict, List, Optional, Union, Tuple
 from .utils import format_number
-import pathlib
+from pathlib import Path
 from .__init__ import __dashboard_url__
 
 
 class SummaryPanel:
     """Holds a summary of the optimization run."""
 
-    def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, runs_dir: str, run_id: str = None):
+    def __init__(
+        self,
+        maximize: bool,
+        metric_name: str,
+        total_steps: int,
+        model: str,
+        runs_dir: str,
+        run_id: str = None,
+        run_name: str = None,
+    ):
         self.maximize = maximize
         self.metric_name = metric_name
-        self.goal = ("Maximizing" if self.maximize else "Minimizing") + f" {self.metric_name}..."
         self.total_input_tokens = 0
         self.total_output_tokens = 0
         self.total_steps = total_steps
         self.model = model
         self.runs_dir = runs_dir
         self.run_id = run_id if run_id is not None else "N/A"
+        self.run_name = run_name if run_name is not None else "N/A"
         self.dashboard_url = "N/A"
+        self.thinking_content = ""
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
             BarColumn(bar_width=20),
             TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
             TextColumn("•"),
-            TextColumn("[bold]{task.completed}/{task.total} Steps"),
+            TextColumn("[bold]{task.completed}/{task.total} Steps "),
             expand=False,
         )
         self.task_id = self.progress.add_task("", total=total_steps)
@@ -40,6 +50,10 @@ class SummaryPanel:
         self.run_id = run_id
         self.set_dashboard_url(run_id=run_id)
 
+    def set_run_name(self, run_name: str):
+        """Set the run name."""
+        self.run_name = run_name
+
     def set_dashboard_url(self, run_id: str):
         """Set the dashboard URL."""
         self.dashboard_url = f"{__dashboard_url__}/runs/{run_id}"
@@ -51,69 +65,100 @@ class SummaryPanel:
     def update_token_counts(self, usage: Dict[str, int]):
         """Update token counts from usage data."""
         if not isinstance(usage, dict) or "input_tokens" not in usage or "output_tokens" not in usage:
-            raise ValueError("Invalid token usage response from API.")
+            raise ValueError("Invalid token usage data received.")
         self.total_input_tokens += usage["input_tokens"]
         self.total_output_tokens += usage["output_tokens"]
 
+    def update_thinking(self, thinking: str):
+        """Update the thinking content."""
+        self.thinking_content = thinking
+
+    def clear_thinking(self):
+        """Clear the thinking content."""
+        self.thinking_content = ""
+
     def get_display(self, final_message: Optional[str] = None) -> Panel:
-        """Create a summary panel with the relevant information."""
-        layout = Layout(name="summary")
-        summary_table = Table(show_header=False, box=None, padding=(0, 1))
+        """Return a Rich panel summarising the current run."""
+        # ───────────────────── summary grid ──────────────────────
+        summary_table = Table.grid(expand=True, padding=(0, 1))
+        summary_table.add_column(ratio=1)
+        summary_table.add_column(justify="right")
+        summary_table.add_row("")
 
+        # Dashboard url
+        summary_table.add_row(f" Dashboard: [underline blue]{self.dashboard_url}[/]")
         summary_table.add_row("")
-        # Goal
+
         if final_message is not None:
-            summary_table.add_row(f"[bold cyan]Result:[/] {final_message}")
-        else:
-            summary_table.add_row(f"[bold cyan]Goal:[/] {self.goal}")
-        summary_table.add_row("")
-        # Model used
-        summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
-        summary_table.add_row("")
-        # Log directory
-        summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.run_id}[/]")
-        summary_table.add_row("")
-        # Dashboard link
-        summary_table.add_row(f"[bold cyan]Dashboard:[/] [blue underline]{self.dashboard_url}[/]")
-        summary_table.add_row("")
-        # Token counts
-        summary_table.add_row(
-            f"[bold cyan]Tokens:[/] ↑[yellow]{format_number(self.total_input_tokens)}[/] ↓[yellow]{format_number(self.total_output_tokens)}[/] = [green]{format_number(self.total_input_tokens + self.total_output_tokens)}[/]"
+            # Add the final message
+            summary_table.add_row(f"[bold cyan] Result:[/] {final_message}", "")
+            summary_table.add_row("")
+
+        # Token info
+        token_info = (
+            f"[bold cyan] {self.model}:[/] "
+            f"↑[yellow]{format_number(self.total_input_tokens)}[/] "
+            f"↓[yellow]{format_number(self.total_output_tokens)}[/] = "
+            f"[green]{format_number(self.total_input_tokens + self.total_output_tokens)} Tokens[/]"
         )
+        summary_table.add_row(token_info)
         summary_table.add_row("")
+
         # Progress bar
         summary_table.add_row(self.progress)
+        summary_table.add_row("")
 
-        # Update layout
-        layout.update(summary_table)
-
-        return Panel(layout, title="[bold]📊 Summary", border_style="magenta", expand=True, padding=(0, 1))
-
-
-class PlanPanel:
-    """Displays the optimization plan with truncation for long plans."""
-
-    def __init__(self):
-        self.plan = ""
+        # Logs url
+        logs_url = Path(self.runs_dir) / self.run_id
+        summary_table.add_row(f" Logs: [underline blue]{logs_url}[/]")
+        summary_table.add_row("")
 
-    def update(self, plan: str):
-        """Update the plan text."""
-        self.plan = plan
+        if final_message is not None:
+            # Don't include the thinking section
+            return Panel(
+                summary_table,
+                title=f"[bold]📊 {'Maximizing' if self.maximize else 'Minimizing'} {self.run_name}",
+                border_style="magenta",
+                expand=True,
+                padding=(0, 1),
+            )
 
-    def clear(self):
-        """Clear the plan text."""
-        self.plan = ""
+        # Include the thinking section
+        layout = Layout(name="summary")
+        layout.split_column(
+            Layout(summary_table, name="main_summary", ratio=1),
+            Layout(
+                Panel(
+                    self.thinking_content or "[dim]No thinking content yet...[/]",
+                    title="[bold]📝 Thinking...",
+                    border_style="cyan",
+                    expand=True,
+                    padding=(0, 1),
+                ),
+                name="thinking_section",
+                ratio=1,
+            ),
+        )
 
-    def get_display(self) -> Panel:
-        """Create a panel displaying the plan with truncation if needed."""
-        return Panel(self.plan, title="[bold]📝 Thinking...", border_style="cyan", expand=True, padding=(0, 1))
+        return Panel(
+            layout,
+            title=f"[bold]📊 {'Maximizing' if self.maximize else 'Minimizing'} {self.run_name}",
+            border_style="magenta",
+            expand=True,
+            padding=(0, 1),
+        )
 
 
 class Node:
     """Represents a node in the solution tree."""
 
     def __init__(
-        self, id: str, parent_id: Union[str, None], code: Union[str, None], metric: Union[float, None], is_buggy: bool
+        self,
+        id: str,
+        parent_id: Union[str, None],
+        code: Union[str, None],
+        metric: Union[float, None],
+        is_buggy: Union[bool, None],
     ):
         self.id = id
         self.parent_id = parent_id
@@ -144,12 +189,15 @@ class MetricTree:
         # Add node to node's parent's children
         if node.parent_id is not None:
             if node.parent_id not in self.nodes:
-                raise ValueError("Could not construct tree: parent node not found.")
+                raise ValueError("Cannot construct optimization tree.")
             self.nodes[node.parent_id].children.append(node)
 
-    def get_draft_nodes(self) -> List[Node]:
-        """Get all draft nodes from the tree."""
-        return [node for node in self.nodes.values() if node.parent_id is None]
+    def get_root_node(self) -> Node:
+        """Get the root node from the tree."""
+        nodes = [node for node in self.nodes.values() if node.parent_id is None]
+        if len(nodes) != 1:
+            raise ValueError("Cannot construct optimization tree.")
+        return nodes[0]
 
     def get_best_node(self) -> Optional[Node]:
         """Get the best node from the tree."""
@@ -157,7 +205,8 @@ class MetricTree:
             node
             for node in self.nodes.values()
             if node.evaluated  # evaluated
-            and not node.is_buggy  # not buggy
+            and node.is_buggy
+            is False  # not buggy => is_buggy can exist in 3 states: None (solution has not yet been evaluated for bugs), True (solution has bug), False (solution does not have a bug)
             and node.metric is not None  # has metric
         ]
         if len(measured_nodes) == 0:
@@ -247,8 +296,8 @@ class MetricTreePanel:
                 append_rec(child, subtree)
 
         tree = Tree("", hide_root=True)
-        for n in self.metric_tree.get_draft_nodes():
-            append_rec(n, tree)
+        root_node = self.metric_tree.get_root_node()
+        append_rec(node=root_node, tree=tree)
 
         return tree
 
@@ -286,7 +335,7 @@ class EvaluationOutputPanel:
 class SolutionPanels:
     """Displays the current and best solutions side by side."""
 
-    def __init__(self, metric_name: str, source_fp: pathlib.Path):
+    def __init__(self, metric_name: str, source_fp: Path):
         # Current solution
         self.current_node = None
         # Best solution
@@ -296,7 +345,7 @@ class SolutionPanels:
         # Determine the lexer for the source file
         self.lexer = self._determine_lexer(source_fp)
 
-    def _determine_lexer(self, source_fp: pathlib.Path) -> str:
+    def _determine_lexer(self, source_fp: Path) -> str:
         """Determine the lexer for the source file."""
         return Syntax.from_path(source_fp).lexer
 
@@ -346,10 +395,7 @@ def create_optimization_layout() -> Layout:
     )
 
     # Split the top section into left and right
-    layout["top_section"].split_row(Layout(name="left_panels", ratio=1), Layout(name="tree", ratio=1))
-
-    # Split the left panels into summary and thinking
-    layout["left_panels"].split_column(Layout(name="summary", ratio=2), Layout(name="plan", ratio=1))
+    layout["top_section"].split_row(Layout(name="summary", ratio=1), Layout(name="tree", ratio=1))
 
     # Split the middle section into left and right
     layout["middle_section"].split_row(Layout(name="current_solution", ratio=1), Layout(name="best_solution", ratio=1))

weco/utils.py

@@ -45,7 +45,7 @@ def determine_default_model(llm_api_keys: Dict[str, Any]) -> str:
         return "gemini-2.5-pro"
     else:
         raise ValueError(
-            "No LLM API keys found in environment. Please set one of the following: OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY."
+            "No LLM API keys found in environment variables. Please set one of the following: OPENAI_API_KEY, ANTHROPIC_API_KEY, or GEMINI_API_KEY based on your model of choice."
         )
 
 
@@ -84,7 +84,7 @@ def write_to_path(fp: pathlib.Path, content: Union[str, Dict[str, Any]], is_json
         elif isinstance(content, str):
             f.write(content)
         else:
-            raise TypeError("Content must be str or Dict[str, Any]")
+            raise TypeError("Error writing to file. Please verify the file path and try again.")
 
 
 # Visualization helper functions
@@ -124,19 +124,54 @@ def smooth_update(
 
 
 # Other helper functions
-def run_evaluation(eval_command: str) -> str:
+DEFAULT_MAX_LINES = 50
+DEFAULT_MAX_CHARS = 5000
+
+
+def truncate_output(output: str, max_lines: int = DEFAULT_MAX_LINES, max_chars: int = DEFAULT_MAX_CHARS) -> str:
+    """Truncate the output to a reasonable size."""
+    lines = output.splitlines()
+
+    # Determine what truncations are needed based on original output
+    lines_truncated = len(lines) > max_lines
+    chars_truncated = len(output) > max_chars
+
+    # Apply truncations to the original output
+    if lines_truncated:
+        output = "\n".join(lines[-max_lines:])
+
+    if chars_truncated:
+        output = output[-max_chars:]
+
+    # Add prefixes for truncations that were applied
+    prefixes = []
+    if lines_truncated:
+        prefixes.append(f"truncated to last {max_lines} lines")
+    if chars_truncated:
+        prefixes.append(f"truncated to last {max_chars} characters")
+
+    if prefixes:
+        prefix_text = ", ".join(prefixes)
+        output = f"... ({prefix_text})\n{output}"
+
+    return output
+
+
+def run_evaluation(eval_command: str, timeout: int | None = None) -> str:
     """Run the evaluation command on the code and return the output."""
 
     # Run the eval command as is
-    result = subprocess.run(eval_command, shell=True, capture_output=True, text=True, check=False)
-
-    # Combine stdout and stderr for complete output
-    output = result.stderr if result.stderr else ""
-    if result.stdout:
-        if len(output) > 0:
-            output += "\n"
-        output += result.stdout
-    return output
+    try:
+        result = subprocess.run(eval_command, shell=True, capture_output=True, text=True, check=False, timeout=timeout)
+        # Combine stdout and stderr for complete output
+        output = result.stderr if result.stderr else ""
+        if result.stdout:
+            if len(output) > 0:
+                output += "\n"
+            output += result.stdout
+        return truncate_output(output)
+    except subprocess.TimeoutExpired:
+        return f"Evaluation timed out after {'an unspecified duration' if timeout is None else f'{timeout} seconds'}."
 
 
 # Update Check Function