weco 0.1.10__py3-none-any.whl → 0.2.0__py3-none-any.whl

weco/panels.py

@@ -0,0 +1,359 @@
+from rich.tree import Tree
+from rich.table import Table
+from rich.progress import BarColumn, Progress, TextColumn
+from rich.layout import Layout
+from rich.panel import Panel
+from rich.syntax import Syntax
+from typing import Dict, List, Optional, Union, Tuple
+from .utils import format_number, truncate_text
+
+
+class SummaryPanel:
+    """Holds a summary of the optimization session."""
+
+    def __init__(self, total_steps: int, model: str, session_id: str = None):
+        self.total_input_tokens = 0
+        self.total_output_tokens = 0
+        self.total_steps = total_steps
+        self.model = model
+        self.session_id = session_id or "unknown"
+        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"),
+            expand=False,
+        )
+        self.task_id = self.progress.add_task("", total=total_steps)
+
+    def set_step(self, step: int):
+        """Set the current step."""
+        self.progress.update(self.task_id, completed=step)
+
+    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.")
+        self.total_input_tokens += usage["input_tokens"]
+        self.total_output_tokens += usage["output_tokens"]
+
+    def get_display(self) -> Panel:
+        """Create a summary panel with the relevant information."""
+        layout = Layout(name="summary")
+        summary_table = Table(show_header=False, box=None, padding=(0, 1))
+        # Log directory
+        runs_dir = f".runs/{self.session_id}"
+        summary_table.add_row(f"[bold cyan]Logs:[/] [blue]{runs_dir}[/]")
+        summary_table.add_row("")
+        # Model used
+        summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
+        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)}[/]"
+        )
+        # Progress bar
+        summary_table.add_row(self.progress)
+
+        # 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 = ""
+        self.max_lines = 3  # Approximate number of lines that can fit in the panel
+
+    def update(self, plan: str):
+        """Update the plan text."""
+        self.plan = plan
+
+    def clear(self):
+        """Clear the plan text."""
+        self.plan = ""
+
+    def get_display(self) -> Panel:
+        """Create a panel displaying the plan with truncation if needed."""
+        display_text = truncate_text(text=self.plan, max_lines=self.max_lines)
+        return Panel(
+            display_text,
+            title="[bold]📝 Thinking...",
+            border_style="cyan",
+            expand=True,
+            padding=(0, 1),
+            subtitle="[cyan]↓ truncated ↓[/]" if len(self.plan) > self.max_lines else None,
+        )
+
+
+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 = id
+        self.parent_id = parent_id
+        self.children: List["Node"] = []
+        self.code = code
+        self.metric = metric
+        self.is_buggy = is_buggy
+        self.evaluated = True
+
+
+class MetricTree:
+    """Manages the tree structure of optimization solutions."""
+
+    def __init__(self, maximize: bool):
+        self.nodes: Dict[str, Node] = {}
+        self.maximize = maximize
+
+    def clear(self):
+        """Clear the tree."""
+        self.nodes = {}
+
+    def add_node(self, node: Node):
+        """Add a node to the tree."""
+        # Add the node to the tree
+        self.nodes[node.id] = node
+
+        # 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.")
+            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_best_node(self) -> Optional[Node]:
+        """Get the best node from the tree."""
+        measured_nodes = [
+            node
+            for node in self.nodes.values()
+            if node.evaluated  # evaluated
+            and not node.is_buggy  # not buggy
+            and node.metric is not None  # has metric
+        ]
+        if len(measured_nodes) == 0:
+            return None
+        if self.maximize:
+            return max(measured_nodes, key=lambda node: node.metric)
+        else:
+            return min(measured_nodes, key=lambda node: node.metric)
+
+
+class MetricTreePanel:
+    """Displays the solution tree with depth limiting."""
+
+    def __init__(self, maximize: bool):
+        self.metric_tree = MetricTree(maximize=maximize)
+
+    def build_metric_tree(self, nodes: List[dict]):
+        """Build the tree from the list of nodes."""
+        # First clear then tree
+        self.metric_tree.clear()
+
+        # Then sort the nodes by step number
+        nodes.sort(key=lambda x: x["step"])
+
+        # Finally build the new tree
+        for node in nodes:
+            self.metric_tree.add_node(
+                Node(
+                    id=node["solution_id"],
+                    parent_id=node["parent_id"],
+                    code=node["code"],
+                    metric=node["metric_value"],
+                    is_buggy=node["is_buggy"],
+                )
+            )
+
+    def set_unevaluated_node(self, node_id: str):
+        """Set the unevaluated node."""
+        self.metric_tree.nodes[node_id].evaluated = False
+
+    def _build_rich_tree(self) -> Tree:
+        """Get a Rich Tree representation of the solution tree using a DFS like traversal."""
+        if len(self.metric_tree.nodes) == 0:
+            return Tree("[bold green]Building 🌳")
+
+        best_node = self.metric_tree.get_best_node()
+
+        def append_rec(node: Node, tree: Tree):
+            if not node.evaluated:
+                # not evaluated
+                color = "yellow"
+                style = None
+                text = "evaluating..."
+            elif node.is_buggy:
+                # buggy node
+                color = "red"
+                style = None
+                text = "bug"
+            else:
+                # evaluated non-buggy node
+                if node.id == best_node.id:
+                    # best node
+                    color = "green"
+                    style = "bold"
+                    text = f"{node.metric:.3f} (best)"
+                elif node.metric is None:
+                    # metric not extracted from evaluated solution
+                    color = "yellow"
+                    style = None
+                    text = "N/A"
+                else:
+                    # evaluated node with metric
+                    color = "green"
+                    style = None
+                    text = f"{node.metric:.3f}"
+
+            s = f"[{f'{style} ' if style is not None else ''}{color}]● {text}"
+            subtree = tree.add(s)
+            for child in node.children:
+                append_rec(child, subtree)
+
+        tree = Tree("🌳")
+        for n in self.metric_tree.get_draft_nodes():
+            append_rec(n, tree)
+
+        return tree
+
+    def get_display(self) -> Panel:
+        """Get a panel displaying the solution tree."""
+        # Make sure the metric tree is built before calling build_rich_tree
+        return Panel(
+            self._build_rich_tree(), title="[bold]🌳 Solution Tree", border_style="green", expand=True, padding=(0, 1)
+        )
+
+
+class EvaluationOutputPanel:
+    """Displays evaluation output with truncation for long outputs."""
+
+    def __init__(self):
+        self.output = ""
+        self.max_lines = 25  # Approximate number of lines that can fit in the panel
+
+    def update(self, output: str) -> None:
+        """Update the evaluation output."""
+        self.output = output
+
+    def clear(self) -> None:
+        """Clear the evaluation output."""
+        self.output = ""
+
+    def get_display(self) -> Panel:
+        """Create a panel displaying the evaluation output with truncation if needed."""
+        display_text = truncate_text(text=self.output, max_lines=self.max_lines)
+
+        title = "[bold]📋 Evaluation Output"
+        if len(self.output) == len(display_text):
+            title += " (truncated)[/]"
+
+        return Panel(display_text, title=title, border_style="red", expand=True, padding=(0, 1))
+
+
+class SolutionPanels:
+    """Displays the current and best solutions side by side."""
+
+    def __init__(self):
+        # Current solution
+        self.current_node = None
+        # Best solution
+        self.best_node = None
+        self.max_lines = 30  # Approximate number of lines that can fit in each panel
+
+    def update(self, current_node: Union[Node, None], best_node: Union[Node, None]):
+        """Update the current and best solutions."""
+        # Update current solution
+        self.current_node = current_node
+        # Update best solution
+        self.best_node = best_node
+
+    def get_display(self, current_step: int) -> Tuple[Panel, Panel]:
+        """Return the current and best solutions as panels."""
+        current_code = self.current_node.code if self.current_node is not None else ""
+        best_code = self.best_node.code if self.best_node is not None else ""
+        best_score = self.best_node.metric if self.best_node is not None else None
+
+        # Determine if code is too long (approximate)
+        current_lines = current_code.count("\n") + 1
+        best_lines = best_code.count("\n") + 1
+
+        # Current solution (without score)
+        current_title = f"[bold]💡 Current Solution (Step {current_step})"
+        current_panel = Panel(
+            Syntax(
+                str(current_code),
+                "python",
+                theme="monokai",
+                line_numbers=True,
+                word_wrap=False,
+                line_range=(0, self.max_lines),  # Only show first max_lines lines
+            ),
+            title=current_title,
+            border_style="yellow",
+            expand=True,
+            padding=(0, 1),
+            subtitle="[yellow]↓ truncated ↓[/]" if current_lines > self.max_lines else None,
+        )
+
+        # Best solution
+        best_title = f"[bold]🏆 Best Solution ([green]Score: {f'{best_score:.4f}' if best_score is not None else 'N/A'}[/])"
+        best_panel = Panel(
+            Syntax(
+                str(best_code),
+                "python",
+                theme="monokai",
+                line_numbers=True,
+                word_wrap=False,
+                line_range=(0, self.max_lines),  # Only show first max_lines lines
+            ),
+            title=best_title,
+            border_style="green",
+            expand=True,
+            padding=(0, 1),
+            subtitle="[yellow]↓ truncated ↓[/]" if best_lines > self.max_lines else None,
+        )
+
+        return current_panel, best_panel
+
+
+def create_optimization_layout() -> Layout:
+    """Create the main layout for the CLI."""
+    layout = Layout()
+
+    # First split into top, middle, and bottom sections
+    layout.split_column(
+        Layout(name="top_section", ratio=3), Layout(name="middle_section", ratio=4), Layout(name="eval_output", ratio=2)
+    )
+
+    # 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))
+
+    # 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))
+
+    return layout
+
+
+def create_end_optimization_layout() -> Layout:
+    """Create the final layout after optimization is complete."""
+    layout = Layout()
+
+    # Create a top section for summary
+    layout.split_column(Layout(name="summary", ratio=1), Layout(name="bottom_section", ratio=3))
+
+    # Split the bottom section into left (best solution) and right ( tree)
+    layout["bottom_section"].split_row(Layout(name="best_solution", ratio=1), Layout(name="tree", ratio=1))
+
+    return layout

weco/utils.py

@@ -1,180 +1,119 @@
-import base64
+from typing import Any, Dict, List, Tuple, Union
+import json
 import os
-import random
-import re
-import string
-from io import BytesIO
-from typing import Dict, Optional, Tuple
-from urllib.parse import urlparse
-
-import requests
-from PIL import Image
-
-
-def is_local_image(maybe_local_image: str) -> bool:
-    """
-    Check if the file is a local image.
-
-    Parameters
-    ----------
-    maybe_local_image : str
-        The file path.
-
-    Returns
-    -------
-    bool
-        True if the file is a local image, False otherwise.
-    """
-    if not os.path.exists(maybe_local_image):  # Check if the file exists
-        return False
-
-    try:  # Check if the file is an image
-        Image.open(maybe_local_image)
-    except IOError:
-        return False
-
-    return True
-
-
-def is_base64_image(maybe_base64: str) -> Tuple[bool, Optional[Dict[str, str]]]:
-    """
-    Check if the image is a base64 encoded image and if so, extract the image information from the encoded data or URL provided.
-
-    Parameters
-    ----------
-    data : str
-        The image data or URL.
-
-    Returns
-    -------
-    Tuple[bool, Optional[Dict[str, str]]]
-    """
-    pattern = r"data:(?P<media_type>[\w/]+);(?P<source_type>\w+),(?P<encoding>.*)"
-    match = re.match(pattern, maybe_base64)
-    if match:
-        return True, match.groupdict()
-
-    return False, None
-
-
-def is_public_url_image(maybe_url_image: str) -> bool:
-    """
-    Check if the string is a publicly accessible URL
-
-    Parameters
-    ----------
-    maybe_url_image : str
-        The URL to check.
-
-    Returns
-    -------
-    bool
-        True if the URL is publicly accessible, False otherwise.
-    """
-    try:
-        # Check if it is a valid URL
-        if not urlparse(maybe_url_image).scheme:
-            return False
-
-        # Check if the URL is publicly accessible
-        response = requests.head(maybe_url_image)
-        if response.status_code != 200:
-            return False
-
-        # Check if the URL is an image
-        content_type = response.headers.get("content-type")
-        if not content_type:
-            return False
-        if not content_type.startswith("image"):
-            return False
-    except Exception:
-        return False
-
-    return True
-
-
-def get_image_size(image: str, source: str) -> float:
-    """
-    Get the size of the image in MB.
-
-    Parameters
-    ----------
-    image : str
-        The image data or URL.
-
-    source : str
-        The source of the image. It can be 'base64', 'url', or 'local'.
-
-    Returns
-    -------
-    float
-        The size of the image in MB.
-
-    Raises
-    ------
-    ValueError
-        If the image is not a valid image.
-    """
-    if source == "base64":
-        _, base64_info = is_base64_image(maybe_base64=image)
-        img_data = base64.b64decode(base64_info["encoding"])
-    elif source == "url":
-        response = requests.get(image)
-        response.raise_for_status()
-        img_data = response.content
-    elif source == "local":
-        with open(image, "rb") as f:
-            img_data = f.read()
+import time
+import subprocess
+from rich.layout import Layout
+from rich.live import Live
+from rich.panel import Panel
+import pathlib
+
+
+# Env/arg helper functions
+def read_api_keys_from_env() -> Dict[str, Any]:
+    """Read API keys from environment variables."""
+    keys = {}
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+    if openai_api_key is not None and len(openai_api_key) > 0:
+        keys["OPENAI_API_KEY"] = openai_api_key
+    if anthropic_api_key is not None and len(anthropic_api_key) > 0:
+        keys["ANTHROPIC_API_KEY"] = anthropic_api_key
+    return keys
+
+
+def read_additional_instructions(additional_instructions: str | None) -> str | None:
+    """Read additional instructions from a file path string or return the string itself."""
+    if additional_instructions is None:
+        return None
+
+    # Try interpreting as a path first
+    potential_path = pathlib.Path(additional_instructions)
+    if potential_path.exists() and potential_path.is_file():
+        return read_from_path(potential_path, is_json=False)  # type: ignore # read_from_path returns str when is_json=False
     else:
-        raise ValueError("Invalid image input")
-
-    img = Image.open(BytesIO(img_data))
-    img_byte_arr = BytesIO()
-    img.save(img_byte_arr, format=img.format)
-    return img_byte_arr.tell() / 1000000  # MB
-
-
-def preprocess_image(image: Image, file_type: str) -> Tuple:
-    """
-    Preprocess the image by converting it to RGB if it has an alpha channel.
-
-    Parameters
-    ----------
-    image : Image
-        The image to preprocess.
-    file_type : str
-        The file type of the image.
-
-    Returns
-    -------
-    Image
-        The preprocessed image.
-    file_type : str
-        The file type of the image.
-    """
-    # Do not rescale or resize. Only do this if latency becomes an issue.
-    # Remove the alpha channel for PNG and WEBP images if it exists.
-    if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
-        image = image.convert("RGB")
-
-    # If the image file type is JPG, convert to JPEG for PIL compatibility.
-    if file_type == "jpg":
-        file_type = "jpeg"
-    return image, file_type
-
-
-def generate_random_base16_code(length: int = 5):
+        # If it's not a valid file path, return the string itself
+        return additional_instructions
+
+
+# File helper functions
+def read_from_path(fp: pathlib.Path, is_json: bool = False) -> Union[str, Dict[str, Any]]:
+    """Read content from a file path, optionally parsing as JSON."""
+    with fp.open("r", encoding="utf-8") as f:
+        if is_json:
+            return json.load(f)
+        return f.read()
+
+
+def write_to_path(fp: pathlib.Path, content: Union[str, Dict[str, Any]], is_json: bool = False) -> None:
+    """Write content to a file path, optionally as JSON."""
+    with fp.open("w", encoding="utf-8") as f:
+        if is_json:
+            json.dump(content, f, indent=4)
+        elif isinstance(content, str):
+            f.write(content)
+        else:
+            raise TypeError("Content must be str or Dict[str, Any]")
+
+
+# Visualization helper functions
+def format_number(n: Union[int, float]) -> str:
+    """Format large numbers with K, M, B, T suffixes for better readability."""
+    if n >= 1e12:
+        return f"{n / 1e12:.1f}T"
+    elif n >= 1e9:
+        return f"{n / 1e9:.1f}B"
+    elif n >= 1e6:
+        return f"{n / 1e6:.1f}M"
+    elif n >= 1e3:
+        return f"{n / 1e3:.1f}K"
+    # Handle potential floats that don't need suffix but might need formatting
+    if isinstance(n, float):
+        # Format floats nicely, avoid excessive precision unless needed
+        return f"{n:.4g}"  # Use general format, up to 4 significant digits
+    return str(n)
+
+
+def smooth_update(
+    live: Live, layout: Layout, sections_to_update: List[Tuple[str, Panel]], transition_delay: float = 0.05
+) -> None:
     """
-    Generate a random base16 code.
+    Update sections of the layout with a small delay between each update for a smoother transition effect.
 
-    Parameters
-    ----------
-    length : int
-        The length of the code.
-
-    Returns
-    -------
-    str
-        The random base16 code.
+    Args:
+        live: The Live display instance
+        layout: The Layout to update
+        sections_to_update: List of (section_name, content) tuples to update
+        transition_delay: Delay in seconds between updates (default: 0.05)
     """
-    return "".join(random.choices(string.hexdigits, k=length))
+    for section, content in sections_to_update:
+        layout[section].update(content)
+        live.refresh()
+        time.sleep(transition_delay)
+
+
+# Other helper functions
+def run_evaluation(eval_command: str) -> 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
+
+
+def truncate_text(text: str, max_lines: int) -> str:
+    """Truncate text to a maximum number of lines."""
+    lines = text.split("\n")
+    if len(lines) > max_lines:
+        display_lines = lines[:max_lines]
+        display_text = "\n".join(display_lines)
+    else:
+        display_text = text
+    return display_text