weco 0.2.22__tar.gz → 0.2.24__tar.gz

{weco-0.2.22 → weco-0.2.24}/.gitignore

@@ -69,8 +69,10 @@ etc/
 # AI generated files
 digest.txt
 .runs/
+CLAUDE.md
 
 *.pyc
 
 # Repomix
+.repomixignore
 repomix-output.*

{weco-0.2.22 → weco-0.2.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.22
+Version: 0.2.24
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -16,6 +16,8 @@ Requires-Dist: requests
 Requires-Dist: rich
 Requires-Dist: packaging
 Requires-Dist: gitingest
+Requires-Dist: fastapi
+Requires-Dist: slowapi
 Provides-Extra: dev
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: build; extra == "dev"
@@ -30,10 +32,11 @@ Dynamic: license-file
 </div>
 
 [![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![PyPI version](https://img.shields.io/pypi/v/weco?label=PyPI%20version&color=f05138&labelColor=555555)](https://badge.fury.io/py/weco)
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
-[![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)
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+[![PyPI Downloads](https://static.pepy.tech/badge/weco?color=4c1)](https://pepy.tech/projects/weco)
+[![arXiv on AIDE](https://img.shields.io/badge/arXiv-AIDE-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg?labelColor=ffffff&color=F17E01)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
 
 `pip install weco`
 
@@ -73,9 +76,9 @@ The `weco` CLI leverages a tree search approach guided by LLMs to iteratively ex
 
     `weco` requires API keys for the LLMs it uses internally. You **must** provide these keys via environment variables:
 
-    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your API key [here](https://platform.openai.com/api-keys))
-    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your API key [here](https://console.anthropic.com/settings/keys))
-    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your OpenAI API key [here](https://platform.openai.com/api-keys))
+    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your Anthropic API key [here](https://console.anthropic.com/settings/keys))
+    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your Gemini API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
 ---
 
@@ -157,6 +160,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
 | `-M, --model`                  | Model identifier for the LLM to use (e.g., `o4-mini`, `claude-sonnet-4-0`).                                                                                                        | `o4-mini` when `OPENAI_API_KEY` is set; `claude-sonnet-4-0` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro` when `GEMINI_API_KEY` is set. | `-M o4-mini`         |
 | `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM.                                                                                             | `None`                                                                                                                                                  | `-i instructions.md` or `-i "Optimize the model for faster inference"`|
 | `-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`             |
 
 ---
 
@@ -247,31 +251,38 @@ Final speedup value = 1.5
 
 Weco will parse this output to extract the numerical value (1.5 in this case) associated with the metric name ('speedup').
 
-## Contributing
+## Supported Models
 
-We welcome your contributions! To get started:
+Weco supports the following LLM models:
 
-1.  **Fork & Clone the Repository:**
+### OpenAI Models
+- `o3`
+- `o3-mini`
+- `o4-mini`
+- `o1-pro`
+- `o1`
+- `gpt-4.1`
+- `gpt-4.1-mini`
+- `gpt-4.1-nano`
+- `gpt-4o`
+- `gpt-4o-mini`
 
-    ```bash
-    git clone https://github.com/WecoAI/weco-cli.git
-    cd weco-cli
-    ```
+### Anthropic Models
+- `claude-opus-4-0`
+- `claude-sonnet-4-0`
+- `claude-3-7-sonnet-latest`
 
-2.  **Install Dependencies:**
-
-    ```bash
-    pip install -e ".[dev]"
-    ```
+### Gemini Models
+- `gemini-2.5-pro`
+- `gemini-2.5-flash`
+- `gemini-2.5-flash-lite`
 
-3.  **Create a Feature Branch:**
+You can specify any of these models using the `-M` or `--model` flag. Ensure you have the corresponding API key set as an environment variable for the model provider you wish to use.
 
-    ```bash
-    git checkout -b feature/your-feature-name
-    ```
+---
 
-4.  **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
+## Contributing
 
-5.  **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.
+We welcome contributions! Please see [contributing.md](contributing.md) for detailed guidelines on how to contribute to this project.
 
 ---

{weco-0.2.22 → weco-0.2.24}/README.md

@@ -6,10 +6,11 @@
 </div>
 
 [![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![PyPI version](https://img.shields.io/pypi/v/weco?label=PyPI%20version&color=f05138&labelColor=555555)](https://badge.fury.io/py/weco)
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
-[![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)
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+[![PyPI Downloads](https://static.pepy.tech/badge/weco?color=4c1)](https://pepy.tech/projects/weco)
+[![arXiv on AIDE](https://img.shields.io/badge/arXiv-AIDE-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg?labelColor=ffffff&color=F17E01)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
 
 `pip install weco`
 
@@ -49,9 +50,9 @@ The `weco` CLI leverages a tree search approach guided by LLMs to iteratively ex
 
     `weco` requires API keys for the LLMs it uses internally. You **must** provide these keys via environment variables:
 
-    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your API key [here](https://platform.openai.com/api-keys))
-    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your API key [here](https://console.anthropic.com/settings/keys))
-    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"` (Create your OpenAI API key [here](https://platform.openai.com/api-keys))
+    - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"` (Create your Anthropic API key [here](https://console.anthropic.com/settings/keys))
+    - **Google:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create your Gemini API key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
 ---
 
@@ -133,6 +134,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
 | `-M, --model`                  | Model identifier for the LLM to use (e.g., `o4-mini`, `claude-sonnet-4-0`).                                                                                                        | `o4-mini` when `OPENAI_API_KEY` is set; `claude-sonnet-4-0` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro` when `GEMINI_API_KEY` is set. | `-M o4-mini`         |
 | `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM.                                                                                             | `None`                                                                                                                                                  | `-i instructions.md` or `-i "Optimize the model for faster inference"`|
 | `-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`             |
 
 ---
 
@@ -223,31 +225,38 @@ Final speedup value = 1.5
 
 Weco will parse this output to extract the numerical value (1.5 in this case) associated with the metric name ('speedup').
 
-## Contributing
+## Supported Models
 
-We welcome your contributions! To get started:
+Weco supports the following LLM models:
 
-1.  **Fork & Clone the Repository:**
+### OpenAI Models
+- `o3`
+- `o3-mini`
+- `o4-mini`
+- `o1-pro`
+- `o1`
+- `gpt-4.1`
+- `gpt-4.1-mini`
+- `gpt-4.1-nano`
+- `gpt-4o`
+- `gpt-4o-mini`
 
-    ```bash
-    git clone https://github.com/WecoAI/weco-cli.git
-    cd weco-cli
-    ```
+### Anthropic Models
+- `claude-opus-4-0`
+- `claude-sonnet-4-0`
+- `claude-3-7-sonnet-latest`
 
-2.  **Install Dependencies:**
-
-    ```bash
-    pip install -e ".[dev]"
-    ```
+### Gemini Models
+- `gemini-2.5-pro`
+- `gemini-2.5-flash`
+- `gemini-2.5-flash-lite`
 
-3.  **Create a Feature Branch:**
+You can specify any of these models using the `-M` or `--model` flag. Ensure you have the corresponding API key set as an environment variable for the model provider you wish to use.
 
-    ```bash
-    git checkout -b feature/your-feature-name
-    ```
+---
 
-4.  **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
+## Contributing
 
-5.  **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.
+We welcome contributions! Please see [contributing.md](contributing.md) for detailed guidelines on how to contribute to this project.
 
 ---

weco-0.2.24/contributing.md

@@ -0,0 +1,23 @@
+# Contributing
+
+We welcome your contributions! To get started:
+
+1.  **Fork & Clone the Repository:**
+    ```bash
+    git clone https://github.com/WecoAI/weco-cli.git
+    cd weco-cli
+    ```
+
+2.  **Install Dependencies:**
+    ```bash
+    pip install -e ".[dev]"
+    ```
+
+3.  **Create a Feature Branch:**
+    ```bash
+    git checkout -b feature/your-feature-name
+    ```
+
+4.  **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
+
+5.  **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.

{weco-0.2.22 → weco-0.2.24}/examples/cuda/README.md

@@ -11,7 +11,7 @@ Install the CLI using `pip`:
 pip install weco
 ```
 
-Set up your API key:
+Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
 ```bash
 export OPENAI_API_KEY="your_key_here"
 ```
@@ -46,3 +46,9 @@ weco run --source optimize.py \
 *   `--additional-instructions guide.md`: Provides guidance to the LLM on the optimization approach.
 
 Weco will iteratively modify `optimize.py`, generating and integrating CUDA C++ code, guided by the evaluation results and the instructions in `guide.md`.
+
+## Next Steps
+
+Now that you've optimized your code with CUDA kernels, try [Triton Optimization](/examples/triton/README.md) for a higher-level GPU programming approach. If you're more interested in [Model Development](/examples/spaceship-titanic/README.md) or [Prompt Engineering](/examples/prompt/README.md), we've got you covered! 
+
+You can check out our [CLI Reference](https://docs.weco.ai/cli/cli-reference) to learn more about what you can do with the tool.

{weco-0.2.22 → weco-0.2.24}/examples/cuda/evaluate.py

@@ -110,6 +110,7 @@ if __name__ == "__main__":
 
     # benchmarking parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -152,6 +153,8 @@ if __name__ == "__main__":
             max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

{weco-0.2.22 → weco-0.2.24}/examples/hello-kernel-world/evaluate.py

@@ -101,6 +101,7 @@ if __name__ == "__main__":
 
     # benchmark parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -130,6 +131,8 @@ if __name__ == "__main__":
         max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size, input_size, args.device)

{weco-0.2.22 → weco-0.2.24}/examples/prompt/README.md

@@ -16,7 +16,7 @@ The experiment runs locally, requires only two short Python files and a prompt g
    pip install weco openai datasets
    ```
 
-3. This example uses `o4-mini` via the OpenAI API by default. Set up your API key:
+3. This example uses `o4-mini` via the OpenAI API by default. Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
    ```bash
    export OPENAI_API_KEY="your_key_here"
    ```
@@ -62,3 +62,10 @@ Weco then mutates the prompt instructions in `optimize.py`, tries again, and gra
 *   The script sends model calls in parallel via `ThreadPoolExecutor`, so network latency is hidden.
 *   Every five completed items, the script logs progress and elapsed time.
 *   The final line `accuracy: value` is the only part Weco needs for guidance.
+
+## Next Steps
+
+Now that you've automated prompt engineering for yourself, check out our guide on [Model Development](/examples/spaceship-titanic/README.md) or [CUDA Kernel Engineering](/examples/cuda/README.md).
+
+You can check out our [CLI Reference](https://docs.weco.ai/cli/cli-reference) to learn more about what you can do with the tool.
+

{weco-0.2.22 → weco-0.2.24}/examples/spaceship-titanic/README.md

@@ -10,7 +10,7 @@ The goal is to improve the model's `accuracy` metric by optimizing the `train.py
     ```bash
     pip install weco pandas numpy scikit-learn torch xgboost lightgbm catboost
     ```
-3. Set up your API key:
+3. Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
     ```bash
     export OPENAI_API_KEY="your_key_here"
     ```
@@ -44,3 +44,7 @@ weco run --source train.py \
 *   `--log-dir .runs/spaceship-titanic`: Specifies the directory where Weco should save logs and results for this run.
 
 Weco will iteratively update the feature engineering or modeling code within `train.py` guided by the evaluation method defined in `evaluate.py`
+
+## Next Steps
+
+With model development covered, you might be curious to see how you can make your AI code run faster, saving you time and more importantly GPU credits. Check out our example on automating kernel engineering in [CUDA](/examples/cuda/README.md) and [Triton](/examples/triton/README.md), or dive into the [CLI Reference](https://docs.weco.ai/cli/cli-reference).

{weco-0.2.22 → weco-0.2.24}/examples/triton/README.md

@@ -12,7 +12,7 @@ Install the CLI using `pip`:
 pip install weco
 ```
 
-Set up your API key:
+Create your OpenAI API key [here](https://platform.openai.com/api-keys), then run:
 ```bash
 export OPENAI_API_KEY="your_key_here"
 ```
@@ -47,3 +47,7 @@ weco run --source optimize.py \
 *   `--additional-instructions "..."`: Provides specific guidance to the LLM. In this case, it directs the model to use Triton for optimization, ensure the numerical difference ("max float diff") between the original and optimized code remains small, and keep the overall code structure consistent.
 
 Weco will iteratively modify `optimize.py`, incorporating Triton kernels, guided by the performance feedback (`speedup`) from the evaluation script and the instructions provided.
+
+## Next Steps
+
+After mastering Triton kernels, explore [CUDA Optimization](/examples/cuda/README.md) for even lower-level GPU programming, or check the [CLI Reference](https://docs.weco.ai/cli/cli-reference) to improve the results you get with Weco.

{weco-0.2.22 → weco-0.2.24}/examples/triton/evaluate.py

@@ -105,6 +105,7 @@ if __name__ == "__main__":
 
     # benchmarking parameters
     n_correctness_trials = 10
+    correctness_tolerance = 1e-5
     n_warmup = 1000
     n_rep = 5000
 
@@ -147,6 +148,8 @@ if __name__ == "__main__":
             max_diff_avg += torch.max(torch.abs(optimized_output - baseline_output))
     max_diff_avg /= n_correctness_trials
     print(f"max float diff between values of baseline and optimized model: {max_diff_avg}")
+    if max_diff_avg > correctness_tolerance:
+        print("invalid solution: max float diff is too high")
 
     # measure performance
     inputs = get_inputs(batch_size=batch_size, seq_len=seq_len, n_embd=n_embd, device="cuda")

{weco-0.2.22 → weco-0.2.24}/pyproject.toml

@@ -8,7 +8,7 @@ name = "weco"
 authors = [{ name = "Weco AI Team", email = "contact@weco.ai" }]
 description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
 readme = "README.md"
-version = "0.2.22"
+version = "0.2.24"
 license = { text = "MIT" }
 requires-python = ">=3.8"
 dependencies = [
@@ -16,6 +16,8 @@ dependencies = [
     "rich",
     "packaging",
     "gitingest",
+    "fastapi",
+    "slowapi",
 ]
 keywords = ["AI", "Code Optimization", "Code Generation"]
 classifiers = [