weco 0.2.15__py3-none-any.whl → 0.2.16__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.16"
 __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

@@ -256,10 +256,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,
@@ -388,7 +385,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
@@ -408,11 +404,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
                     )
 
                     # Update the step of the progress bar
@@ -493,7 +485,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 +501,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"])
@@ -575,7 +566,11 @@ def main() -> None:
             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.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.15
+Version: 0.2.16
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -20,7 +20,7 @@ Requires-Dist: build; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
 Dynamic: license-file
 
-# Weco: The Evaluation-Driven AI Code Optimizer
+# Weco: The AI Research Engineer
 
 [![Python](https://img.shields.io/badge/Python-3.12.0-blue)](https://www.python.org)
 [![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
@@ -30,11 +30,11 @@ Weco systematically optimizes your code, guided directly by your evaluation metr
 
 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 +42,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 +101,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)**, 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

weco-0.2.16.dist-info/RECORD

@@ -0,0 +1,12 @@
+weco/__init__.py,sha256=HyOG-766A8if3MM5YvAM6CMJobEib3oM54UAJMk9GFA,426
+weco/api.py,sha256=0OU1hEhN7sbIZ1zj8TeeB0tMaxXk6n6qw82FmcdK0ec,3111
+weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
+weco/cli.py,sha256=DPObmIes3tRAVswBGYOx5G_8zIjeIJFufOGnFt__XwE,28869
+weco/panels.py,sha256=8DoTQC-epGpGjn-xDBcqelC5BKaX7JXnrJ97LInEbRU,13561
+weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
+weco-0.2.16.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.16.dist-info/METADATA,sha256=CFe6w8pRnMh-zDuyebtqIOFfIuf6JEJFqgVdUPSTO6s,10749
+weco-0.2.16.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+weco-0.2.16.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.16.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.16.dist-info/RECORD,,

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

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (79.0.1)
+Generator: setuptools (80.0.0)
 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,,