weco 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

weco/__init__.py

@@ -1,4 +1,4 @@
 # DO NOT EDIT
-__pkg_version__ = "0.2.7"
+__pkg_version__ = "0.2.9"
 __api_version__ = "v1"
 __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"

weco/api.py

@@ -6,14 +6,9 @@ import sys
 
 
 def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
-    """Extract and display error messages from API responses."""
-    try:
-        error_data = e.response.json()
-        error_message = error_data.get("detail", str(e))
-        console.print(f"[bold red]Server Error:[/] {error_message}")
-    except Exception:
-        # If we can't parse the JSON, just show the original error
-        console.print(f"[bold red]Server Error:[/] {str(e)}")
+    """Extract and display error messages from API responses in a structured format."""
+    error_message = str(e)  # Default message
+    console.print(f"[bold red]Error:[/] {error_message}")
     sys.exit(1)
 
 

weco/cli.py

@@ -36,7 +36,7 @@ def main() -> None:
     parser = argparse.ArgumentParser(
         description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
     )
-    parser.add_argument("--source", type=str, required=True, help="Path to the Python source code (e.g. optimize.py)")
+    parser.add_argument("--source", type=str, required=True, help="Path to the source code (e.g. optimize.py)")
     parser.add_argument(
         "--eval-command", type=str, required=True, help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1')"
     )
@@ -57,6 +57,11 @@ def main() -> None:
         type=str,
         help="Description of additional instruction or path to a file containing additional instructions",
     )
+    parser.add_argument(
+        "--preserve-source",
+        action="store_true",
+        help="If set, do not overwrite the original source file; only save modified versions in the runs directory",
+    )
     args = parser.parse_args()
 
     try:
@@ -73,22 +78,23 @@ def main() -> None:
                 "debug_prob": 0.5,
                 "max_debug_depth": max(1, math.ceil(0.1 * steps)),  # 10% of steps
             }
+            # Read API keys
+            api_keys = read_api_keys_from_env()
+            # API request timeout
+            timeout = 800
+
             # Read additional instructions
             additional_instructions = read_additional_instructions(additional_instructions=args.additional_instructions)
             # Read source code
             source_fp = pathlib.Path(args.source)
             source_code = read_from_path(fp=source_fp, is_json=False)
-            # Read API keys
-            api_keys = read_api_keys_from_env()
-            # API request timeout
-            timeout = 800
 
         # Initialize panels
         summary_panel = SummaryPanel(
             maximize=maximize, metric_name=metric_name, total_steps=steps, model=args.model, runs_dir=args.log_dir
         )
         plan_panel = PlanPanel()
-        solution_panels = SolutionPanels(metric_name=metric_name)
+        solution_panels = SolutionPanels(metric_name=metric_name, source_fp=source_fp)
         eval_output_panel = EvaluationOutputPanel()
         tree_panel = MetricTreePanel(maximize=maximize)
         layout = create_optimization_layout()
@@ -118,13 +124,14 @@ def main() -> None:
             runs_dir = pathlib.Path(args.log_dir) / session_id
             runs_dir.mkdir(parents=True, exist_ok=True)
 
-            # Save the original code (.runs/<session-id>/original.py)
-            runs_copy_source_fp = runs_dir / "original.py"
+            # Save the original code (.runs/<session-id>/original.<extension>)
+            runs_copy_source_fp = runs_dir / f"original.{source_fp.suffix}"
             write_to_path(fp=runs_copy_source_fp, content=source_code)
 
             # Write the code string to the source file path
             # Do this after the original code is saved
-            write_to_path(fp=source_fp, content=session_response["code"])
+            if not args.preserve_source:
+                write_to_path(fp=source_fp, content=session_response["code"])
 
             # Update the panels with the initial solution
             # Add session id now that we have it
@@ -191,20 +198,25 @@ def main() -> None:
             )
 
             for step in range(1, steps):
+                # Re-read instructions from the original source (file path or string) BEFORE each suggest call
+                current_additional_instructions = read_additional_instructions(
+                    additional_instructions=args.additional_instructions
+                )
                 # Evaluate the current output and get the next solution
                 eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
                     console=console,
                     session_id=session_id,
                     execution_output=term_out,
-                    additional_instructions=additional_instructions,
+                    additional_instructions=current_additional_instructions,
                     api_keys=api_keys,
                     timeout=timeout,
                 )
-                # Save next solution (.runs/<session-id>/step_<step>.py)
-                write_to_path(fp=runs_dir / f"step_{step}.py", content=eval_and_next_solution_response["code"])
+                # Save next solution (.runs/<session-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"])
+                if not args.preserve_source:
+                    write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
 
                 # Get the optimization session status for
                 # the best solution, its score, and the history to plot the tree
@@ -283,12 +295,16 @@ def main() -> None:
                     transition_delay=0.1,  # Slightly longer delay for evaluation results
                 )
 
+            # Re-read instructions before the final feedback step
+            current_additional_instructions = read_additional_instructions(
+                additional_instructions=args.additional_instructions
+            )
             # Ensure we pass evaluation results for the last step's generated solution
             eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
                 console=console,
                 session_id=session_id,
                 execution_output=term_out,
-                additional_instructions=additional_instructions,
+                additional_instructions=current_additional_instructions,
                 api_keys=api_keys,
                 timeout=timeout,
             )
@@ -351,11 +367,12 @@ def main() -> None:
                 )
                 best_solution_content = f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
 
-            # Save best solution to .runs/<session-id>/best.py
-            write_to_path(fp=runs_dir / "best.py", content=best_solution_content)
+            # Save best solution to .runs/<session-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)
+            if not args.preserve_source:
+                write_to_path(fp=source_fp, content=best_solution_content)
 
         console.print(end_optimization_layout)
 

weco/panels.py

@@ -6,6 +6,7 @@ from rich.panel import Panel
 from rich.syntax import Syntax
 from typing import Dict, List, Optional, Union, Tuple
 from .utils import format_number
+import pathlib
 
 
 class SummaryPanel:
@@ -46,6 +47,8 @@ class SummaryPanel:
         """Create a summary panel with the relevant information."""
         layout = Layout(name="summary")
         summary_table = Table(show_header=False, box=None, padding=(0, 1))
+
+        summary_table.add_row("")
         # Goal
         if final_message is not None:
             summary_table.add_row(f"[bold cyan]Result:[/] {final_message}")
@@ -256,13 +259,19 @@ class EvaluationOutputPanel:
 class SolutionPanels:
     """Displays the current and best solutions side by side."""
 
-    def __init__(self, metric_name: str):
+    def __init__(self, metric_name: str, source_fp: pathlib.Path):
         # Current solution
         self.current_node = None
         # Best solution
         self.best_node = None
         # Metric name
         self.metric_name = metric_name.capitalize()
+        # Determine the lexer for the source file
+        self.lexer = self._determine_lexer(source_fp)
+
+    def _determine_lexer(self, source_fp: pathlib.Path) -> str:
+        """Determine the lexer for the source file."""
+        return Syntax.from_path(source_fp).lexer
 
     def update(self, current_node: Union[Node, None], best_node: Union[Node, None]):
         """Update the current and best solutions."""
@@ -280,7 +289,7 @@ class SolutionPanels:
         # 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),
+            Syntax(str(current_code), self.lexer, theme="monokai", line_numbers=True, word_wrap=False),
             title=current_title,
             border_style="yellow",
             expand=True,
@@ -290,7 +299,7 @@ class SolutionPanels:
         # Best solution
         best_title = f"[bold]🏆 Best Solution ([green]{self.metric_name}: {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),
+            Syntax(str(best_code), self.lexer, theme="monokai", line_numbers=True, word_wrap=False),
             title=best_title,
             border_style="green",
             expand=True,

{weco-0.2.7.dist-info → weco-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.7
+Version: 0.2.9
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -9,7 +9,7 @@ Keywords: AI,Code Optimization,Code Generation
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Classifier: License :: OSI Approved :: MIT License
-Requires-Python: >=3.12
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
@@ -20,13 +20,19 @@ Requires-Dist: build; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
 Dynamic: license-file
 
-# Weco CLI – Code Optimizer for Machine Learning Engineers
+# Weco: The Evaluation-Driven AI Code Optimizer
 
 [![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
-[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
 [![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
+[![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
 
-`weco` is a command-line interface for interacting with Weco AI's code optimizer, powered by [AI-Driven Exploration](https://arxiv.org/abs/2502.13138). It helps you automate the improvement of your code for tasks like GPU kernel optimization, feature engineering, model development, and prompt engineering.
+Weco systematically optimizes your code, guided directly by your evaluation metrics.
+
+Example applications include:
+
+- **GPU Kernel Optimization**: Reimplement PyTorch functions using CUDA, Triton or Metal, optimizing for `latency`, `throughput`, or `memory_bandwidth`.
+- **Model Development**: Tune feature transformations or architectures, optimizing for `validation_accuracy`, `AUC`, or `Sharpe Ratio`.
+- **Prompt Engineering**: Refine prompts for LLMs, optimizing for  `win_rate`, `relevance`, or `format_adherence`
 
 https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
 
@@ -40,37 +46,6 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
 
 ---
 
-## Example Use Cases
-
-Here's how `weco` can be applied to common ML engineering tasks:
-
-*   **GPU Kernel Optimization:**
-    *   **Goal:** Improve the speed or efficiency of low-level GPU code.
-    *   **How:** `weco` iteratively refines CUDA, Triton, Metal, or other kernel code specified in your `--source` file.
-    *   **`--eval-command`:** Typically runs a script that compiles the kernel, executes it, and benchmarks performance (e.g., latency, throughput).
-    *   **`--metric`:** Examples include `latency`, `throughput`, `TFLOPS`, `memory_bandwidth`. Optimize to `minimize` latency or `maximize` throughput.
-
-*   **Feature Engineering:**
-    *   **Goal:** Discover better data transformations or feature combinations for your machine learning models.
-    *   **How:** `weco` explores different processing steps or parameters within your feature transformation code (`--source`).
-    *   **`--eval-command`:** Executes a script that applies the features, trains/validates a model using those features, and prints a performance score.
-    *   **`--metric`:** Examples include `accuracy`, `AUC`, `F1-score`, `validation_loss`. Usually optimized to `maximize` accuracy/AUC/F1 or `minimize` loss.
-
-*   **Model Development:**
-    *   **Goal:** Tune hyperparameters or experiment with small architectural changes directly within your model's code.
-    *   **How:** `weco` modifies hyperparameter values (like learning rate, layer sizes if defined in the code) or structural elements in your model definition (`--source`).
-    *   **`--eval-command`:** Runs your model training and evaluation script, printing the key performance indicator.
-    *   **`--metric`:** Examples include `validation_accuracy`, `test_loss`, `inference_time`, `perplexity`. Optimize according to the metric's nature (e.g., `maximize` accuracy, `minimize` loss).
-
-*   **Prompt Engineering:**
-    *   **Goal:** Refine prompts used within larger systems (e.g., for LLM interactions) to achieve better or more consistent outputs.
-    *   **How:** `weco` modifies prompt templates, examples, or instructions stored in the `--source` file.
-    *   **`--eval-command`:** Executes a script that uses the prompt, generates an output, evaluates that output against desired criteria (e.g., using another LLM, checking for keywords, format validation), and prints a score.
-    *   **`--metric`:** Examples include `quality_score`, `relevance`, `task_success_rate`, `format_adherence`. Usually optimized to `maximize`.
-
----
-
-
 ## Setup
 
 1.  **Install the Package:**
@@ -97,13 +72,20 @@ Here's how `weco` can be applied to common ML engineering tasks:
 
 ---
 
-### Examples
+### Example: Optimizing Simple PyTorch Operations
+
+This basic example shows how to optimize a simple PyTorch function for speedup.
 
-**Example 1: Optimizing PyTorch simple operations**
+For more advanced examples, including **[Metal/MLX](/examples/metal/README.md), [Triton](/examples/triton/README.md), [CUDA kernel optimization](/examples/cuda/README.md)**, and **[ML model optimization](/examples/spaceship-titanic/README.md)**, please see the `README.md` files within the corresponding subdirectories under the [`examples/`](./examples/) folder.
 
 ```bash
+# Navigate to the example directory
 cd examples/hello-kernel-world
-pip install torch 
+
+# Install dependencies
+pip install torch
+
+# Run Weco
 weco --source optimize.py \
      --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
      --metric speedup \
@@ -113,96 +95,7 @@ weco --source optimize.py \
      --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
 ```
 
-Note that if you have an NVIDIA gpu, change the device to `cuda`. If you are running this on Apple Silicon, set it to `mps`.
-
-**Example 2: Optimizing MLX operations with instructions from a file**
-
-Lets optimize a 2D convolution operation in [`mlx`](https://github.com/ml-explore/mlx) using [Metal](https://developer.apple.com/documentation/metal/). Sometimes, additional context or instructions are too complex for a single command-line string. You can provide a path to a file containing these instructions.
-
-```bash
-cd examples/metal
-pip install mlx
-weco --source optimize.py \
-     --eval-command "python evaluate.py --solution-path optimize.py" \
-     --metric speedup \
-     --maximize true \
-     --steps 30 \
-     --model gemini-2.5-pro-exp-03-25 \
-     --additional-instructions examples.rst
-```
-
-**Example 3: Level Agnostic Optimization: Causal Self Attention with Triton & CUDA**
-
-Given how useful causal multihead self attention is to transformers, we've seen its wide adoption across ML engineering and AI research. Its great to keep things at a high-level (in PyTorch) when doing research, but when moving to production you often need to write highly customized low-level kernels to make things run as fast as they can. The `weco` CLI can optimize kernels across a variety of different abstraction levels and frameworks. Example 2 uses Metal but lets explore two more frameworks:
-
-1. [Triton](https://github.com/triton-lang/triton)
-    ```bash
-   cd examples/triton
-   pip install torch triton
-   weco --source optimize.py \
-        --eval-command "python evaluate.py --solution-path optimize.py" \
-        --metric speedup \
-        --maximize true \
-        --steps 30 \
-        --model gemini-2.5-pro-exp-03-25 \
-        --additional-instructions "Use triton to optimize the code while ensuring a small max float diff. Maintain the same code format."
-   ```
-
-2. [CUDA](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html)
-   ```bash
-   cd examples/cuda
-   pip install torch
-   weco --source optimize.py \
-        --eval-command "python evaluate.py --solution-path optimize.py" \
-        --metric speedup \
-        --maximize true \
-        --steps 30 \
-        --model gemini-2.5-pro-exp-03-25 \
-        --additional-instructions guide.md
-   ```
-
-**Example 4: Optimizing a Classification Model**
-
-This example demonstrates optimizing a script for a Kaggle competition ([Spaceship Titanic](https://www.kaggle.com/competitions/spaceship-titanic/overview)) to improve classification accuracy. The additional instructions are provided via a separate file (`examples/spaceship-titanic/README.md`).
-
-First, install the requirements for the example environment:
-```bash
-pip install -r examples/spaceship-titanic/requirements-test.txt
-```
-And run utility function once to prepare the dataset
-```bash
-python examples/spaceship-titanic/utils.py
-```
-
-You should see the following structure at `examples/spaceship-titanic`. You need to prepare the kaggle credentials for downloading the dataset.
-```
-.
-├── baseline.py
-├── evaluate.py
-├── optimize.py
-├── private
-│   └── test.csv
-├── public
-│   ├── sample_submission.csv
-│   ├── test.csv
-│   └── train.csv
-├── README.md
-├── requirements-test.txt
-└── utils.py
-```
-
-Then, execute the optimization command:
-```bash
-weco --source examples/spaceship-titanic/optimize.py \
-     --eval-command "python examples/spaceship-titanic/optimize.py && python examples/spaceship-titanic/evaluate.py" \
-     --metric accuracy \
-     --maximize true \
-     --steps 10 \
-     --model gemini-2.5-pro-exp-03-25 \
-     --additional-instructions examples/spaceship-titanic/README.md
-```
-
-*The [baseline.py](examples/spaceship-titanic/baseline.py) is provided as a start point for optimization*
+**Note:** If you have an NVIDIA GPU, change the device in the `--eval-command` to `cuda`. If you are running this on Apple Silicon, set it to `mps`.
 
 ---
 
@@ -215,9 +108,10 @@ weco --source examples/spaceship-titanic/optimize.py \
 | `--metric`                  | The name of the metric you want to optimize (e.g., 'accuracy', 'speedup', 'loss'). This metric name should match what's printed by your `--eval-command`.                | Yes      |
 | `--maximize`                | Whether to maximize (`true`) or minimize (`false`) the metric.                                                                                                           | Yes      |
 | `--steps`                   | Number of optimization steps (LLM iterations) to run.                                                                                                                    | Yes      |
-| `--model`                   | Model identifier for the LLM to use (e.g., `gpt-4o`, `claude-3.5-sonnet`). Recommended models to try include `o3-mini`, `claude-3-haiku`, and `gemini-2.5-pro-exp-03-25`.| Yes      |
+| `--model`                   | Model identifier for the LLM to use (e.g., `gpt-4o`, `claude-3.7-sonnet`). Recommended models to try include `o4-mini`, and `gemini-2.5-pro-exp-03-25`.| Yes      |
 | `--additional-instructions` | (Optional) Natural language description of specific instructions OR path to a file containing detailed instructions to guide the LLM.                                    | No       |
 | `--log-dir`                 | (Optional) Path to the directory to log intermediate steps and final optimization result. Defaults to `.runs/`.                                                          | No       |
+| `--preserve-source`         | (Optional) If set, do not overwrite the original `--source` file. Modifications and the best solution will still be saved in the `--log-dir`.                                | No       |
 
 ---
 

weco-0.2.9.dist-info/RECORD

@@ -0,0 +1,11 @@
+weco/__init__.py,sha256=8vfBL3OdlOveq9nYktaI3JQLXAeR7Pnwz8TRZ3xu0nA,124
+weco/api.py,sha256=89lB2572jApAxkA0DDppDnJKBwvZTa3kH9jFpC0LFDQ,3313
+weco/cli.py,sha256=iTecBQ0gAqyQixj7tNU9S6f9NW7TMSGZtH8nw-85_Oc,18276
+weco/panels.py,sha256=R_df-VAbWyLoqCA9A6UzbIGZ9sm2IgJO4idnyjmrHQk,12701
+weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
+weco-0.2.9.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.9.dist-info/METADATA,sha256=Fq37xE2H8Ia3brOv6AcjsleqDeDL_U6XiV3bJHZOhCY,9378
+weco-0.2.9.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+weco-0.2.9.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.9.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.9.dist-info/RECORD,,

{weco-0.2.7.dist-info → weco-0.2.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

weco-0.2.7.dist-info/RECORD

@@ -1,11 +0,0 @@
-weco/__init__.py,sha256=6ZYuD51wx6bytYdLFMvzihYrpBlSxiFbmuiEl0z_AFo,124
-weco/api.py,sha256=8rIf2Fy3tN6GW7BG1CaggtfE9pW56I1erzwLCgawcVE,3511
-weco/cli.py,sha256=h8FevpztBob7OziDTIKh4y9CSnehkHJp2ydts6V6DhM,17317
-weco/panels.py,sha256=HHWmrnc2EJBJ8AEHb8mxlq30m-3q0j_4axc0r17sTnE,12349
-weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
-weco-0.2.7.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.7.dist-info/METADATA,sha256=AlBfVky2jZ6C2MnB8k_WkSMyucTeTE1Ias3y3_4bVsE,14577
-weco-0.2.7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-weco-0.2.7.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.7.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.7.dist-info/RECORD,,