weco 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

weco/__init__.py

@@ -1,7 +1,7 @@
 import os
 
 # DO NOT EDIT
-__pkg_version__ = "0.2.15"
+__pkg_version__ = "0.2.17"
 __api_version__ = "v1"
 
 __base_url__ = f"https://api.weco.ai/{__api_version__}"

weco/api.py

@@ -28,33 +28,28 @@ def start_optimization_session(
 ) -> Dict[str, Any]:
     """Start the optimization session."""
     with console.status("[bold green]Starting Optimization..."):
-        try:
-            # __base_url__ already contains /v1
-            response = requests.post(
-                f"{__base_url__}/sessions",  # Path is relative to base_url
-                json={
-                    "source_code": source_code,
-                    "additional_instructions": additional_instructions,
-                    "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
-                    "optimizer": {
-                        "steps": steps,
-                        "code_generator": code_generator_config,
-                        "evaluator": evaluator_config,
-                        "search_policy": search_policy_config,
-                    },
-                    "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
+        response = requests.post(
+            f"{__base_url__}/sessions",  # Path is relative to base_url
+            json={
+                "source_code": source_code,
+                "additional_instructions": additional_instructions,
+                "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
+                "optimizer": {
+                    "steps": steps,
+                    "code_generator": code_generator_config,
+                    "evaluator": evaluator_config,
+                    "search_policy": search_policy_config,
                 },
-                headers=auth_headers,  # Add headers
-                timeout=timeout,
-            )
-            response.raise_for_status()
-            return response.json()
-        except requests.exceptions.HTTPError as e:
-            handle_api_error(e=e, console=console)
+                "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
+            },
+            headers=auth_headers,  # Add headers
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        return response.json()
 
 
 def evaluate_feedback_then_suggest_next_solution(
-    console: rich.console.Console,
     session_id: str,
     execution_output: str,
     additional_instructions: str = None,
@@ -63,41 +58,29 @@ def evaluate_feedback_then_suggest_next_solution(
     timeout: int = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
-    try:
-        # __base_url__ already contains /v1
-        response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
-            json={
-                "execution_output": execution_output,
-                "additional_instructions": additional_instructions,
-                "metadata": {**api_keys},
-            },
-            headers=auth_headers,  # Add headers
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e=e, console=console)
+    response = requests.post(
+        f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
+        json={
+            "execution_output": execution_output,
+            "additional_instructions": additional_instructions,
+            "metadata": {**api_keys},
+        },
+        headers=auth_headers,  # Add headers
+        timeout=timeout,
+    )
+    response.raise_for_status()
+    return response.json()
 
 
 def get_optimization_session_status(
-    console: rich.console.Console,
-    session_id: str,
-    include_history: bool = False,
-    auth_headers: dict = {},
-    timeout: int = 800,  # Add auth_headers
+    session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
 ) -> Dict[str, Any]:
     """Get the current status of the optimization session."""
-    try:
-        # __base_url__ already contains /v1
-        response = requests.get(
-            f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
-            params={"include_history": include_history},
-            headers=auth_headers,  # Add headers
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e=e, console=console)
+    response = requests.get(
+        f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
+        params={"include_history": include_history},
+        headers=auth_headers,
+        timeout=timeout,
+    )
+    response.raise_for_status()
+    return response.json()

weco/cli.py

@@ -194,11 +194,6 @@ def main() -> None:
         type=str,
         help="Description of additional instruction or path to a file containing additional instructions",
     )
-    run_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",
-    )
 
     # --- Logout Command ---
     _ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
@@ -256,10 +251,7 @@ def main() -> None:
             maximize = args.maximize == "true"
             steps = args.steps
             code_generator_config = {"model": args.model}
-            evaluator_config = {
-                "model": args.model,
-                "include_analysis": False,  # NOTE: False for now
-            }
+            evaluator_config = {"model": args.model, "include_analysis": True}
             search_policy_config = {
                 "num_drafts": max(1, math.ceil(0.15 * steps)),
                 "debug_prob": 0.5,
@@ -313,9 +305,8 @@ def main() -> None:
                 # Write the initial code string to the logs
                 write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=session_response["code"])
 
-                # Write the initial code string to the source file path (if not preserving)
-                if not args.preserve_source:
-                    write_to_path(fp=source_fp, content=session_response["code"])
+                # Write the initial code string to the source file path
+                write_to_path(fp=source_fp, content=session_response["code"])
 
                 # Update the panels with the initial solution
                 summary_panel.set_session_id(session_id=session_id)  # Add session id now that we have it
@@ -388,7 +379,6 @@ def main() -> None:
 
                     # Send feedback and get next suggestion
                     eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
-                        console=console,
                         session_id=session_id,
                         execution_output=term_out,
                         additional_instructions=current_additional_instructions,  # Pass current instructions
@@ -402,17 +392,12 @@ def main() -> None:
                     )
 
                     # Write the next solution to the source file
-                    if not args.preserve_source:
-                        write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
+                    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
                     status_response = get_optimization_session_status(
-                        console=console,
-                        session_id=session_id,
-                        include_history=True,
-                        timeout=timeout,
-                        auth_headers=auth_headers,
+                        session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                     )
 
                     # Update the step of the progress bar
@@ -493,7 +478,6 @@ def main() -> None:
 
                 # 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=current_additional_instructions,
@@ -510,7 +494,7 @@ def main() -> None:
                 # Get the optimization session status for
                 # the best solution, its score, and the history to plot the tree
                 status_response = get_optimization_session_status(
-                    console=console, session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                 )
                 # Build the metric tree
                 tree_panel.build_metric_tree(nodes=status_response["history"])
@@ -569,13 +553,16 @@ def main() -> None:
                 write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
 
                 # write the best solution to the source file
-                if not args.preserve_source:
-                    write_to_path(fp=source_fp, content=best_solution_content)
+                write_to_path(fp=source_fp, content=best_solution_content)
 
             console.print(end_optimization_layout)
 
         except Exception as e:
-            console.print(Panel(f"[bold red]Error: {str(e)}", title="[bold red]Error", border_style="red"))
+            try:
+                error_message = e.response.json()["detail"]
+            except Exception:
+                error_message = str(e)
+            console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Error", border_style="red"))
             # Print traceback for debugging
-            console.print_exception(show_locals=True)
+            # console.print_exception(show_locals=False)
             sys.exit(1)

weco/panels.py

@@ -121,6 +121,7 @@ class Node:
         self.metric = metric
         self.is_buggy = is_buggy
         self.evaluated = True
+        self.name = ""
 
 
 class MetricTree:
@@ -181,16 +182,17 @@ class MetricTreePanel:
         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"],
-                )
+        for i, node in enumerate(nodes):
+            node = Node(
+                id=node["solution_id"],
+                parent_id=node["parent_id"],
+                code=node["code"],
+                metric=node["metric_value"],
+                is_buggy=node["is_buggy"],
             )
+            if i == 0:
+                node.name = "baseline"
+            self.metric_tree.add_node(node)
 
     def set_unevaluated_node(self, node_id: str):
         """Set the unevaluated node."""
@@ -232,12 +234,15 @@ class MetricTreePanel:
                     style = None
                     text = f"{node.metric:.3f}"
 
+                # add the node name info
+                text = f"{node.name} {text}".strip()
+
             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("🌳")
+        tree = Tree("", hide_root=True)
         for n in self.metric_tree.get_draft_nodes():
             append_rec(n, tree)
 

{weco-0.2.15.dist-info → weco-0.2.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.15
+Version: 0.2.17
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -20,21 +20,30 @@ Requires-Dist: build; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
 Dynamic: license-file
 
-# Weco: The Evaluation-Driven AI Code Optimizer
+<div align="center">
 
-[![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
+# Weco: The AI Code Optimizer
+
+[![Python](https://img.shields.io/badge/Python-3.8.0+-blue)](https://www.python.org)
+[![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)
 
+<code>pip install weco</code>
+
+</div>
+
+---
+
 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`.
+- **GPU Kernel Optimization**: Reimplement PyTorch functions using CUDA or Triton 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
+![image](assets/example-optimization.gif)
 
 ---
 
@@ -42,7 +51,7 @@ https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
 
 The `weco` CLI leverages a tree search approach guided by Large Language Models (LLMs) to iteratively explore and refine your code. It automatically applies changes, runs your evaluation script, parses the results, and proposes further improvements based on the specified goal.
 
-[image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
+![image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
 
 ---
 
@@ -101,7 +110,7 @@ This command starts the optimization process.
 
 This basic example shows how to optimize a simple PyTorch function for speedup.
 
-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.
+For more advanced examples, including [Triton](/examples/triton/README.md), [CUDA kernel optimization](/examples/cuda/README.md), [ML model optimization](/examples/spaceship-titanic/README.md), and [prompt engineering for math problems](https://github.com/WecoAI/weco-cli/tree/main/examples/prompt), please see the `README.md` files within the corresponding subdirectories under the [`examples/`](./examples/) folder.
 
 ```bash
 # Navigate to the example directory
@@ -136,7 +145,6 @@ weco run --source optimize.py \
 | `--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      |
 | `--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.17.dist-info/RECORD

@@ -0,0 +1,12 @@
+weco/__init__.py,sha256=aMh5ZK_EMG8kGfy59ah_5gRVfWeEXgBhcfReKNsqNmQ,426
+weco/api.py,sha256=0OU1hEhN7sbIZ1zj8TeeB0tMaxXk6n6qw82FmcdK0ec,3111
+weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
+weco/cli.py,sha256=j8EyHVIIl2zNAjfVUoOFtJBZbDV69LrfnFA2WlDgbao,28488
+weco/panels.py,sha256=8DoTQC-epGpGjn-xDBcqelC5BKaX7JXnrJ97LInEbRU,13561
+weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
+weco-0.2.17.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.17.dist-info/METADATA,sha256=G2unWMKcPFan2r5a1kGlcizCQr1nVayJBVVLb-e-9fE,10795
+weco-0.2.17.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+weco-0.2.17.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.17.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.17.dist-info/RECORD,,

{weco-0.2.15.dist-info → weco-0.2.17.dist-info}/WHEEL

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

weco-0.2.15.dist-info/RECORD

@@ -1,12 +0,0 @@
-weco/__init__.py,sha256=2P-Z80aRiNeL8hyXD4ExZxizUgLb_Tj_NOydwoqTZ_k,426
-weco/api.py,sha256=z2DCe0kQJaaBHo-Vml52GRc9nwQBnrKATXoY3UfPojw,3842
-weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
-weco/cli.py,sha256=rcntl-hFQorCNSBX1yaNoL-4eBJFSkN5Vo35b6-W4UU,28993
-weco/panels.py,sha256=gB4rZbCvqzewUCBcILvyyU4fnOQLwFgHCGmtn-ZlgSo,13385
-weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
-weco-0.2.15.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.15.dist-info/METADATA,sha256=B97o-KOP1o2sUVHqSz1maZGy-iysBKBbvq16QIyVE68,10851
-weco-0.2.15.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-weco-0.2.15.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.15.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.15.dist-info/RECORD,,