weco 0.2.14__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.14"
+__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,23 +256,20 @@ 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,
                 "max_debug_depth": max(1, math.ceil(0.1 * steps)),
             }
+            # API request timeout
+            timeout = 800
             # Read additional instructions
             additional_instructions = read_additional_instructions(additional_instructions=args.additional_instructions)
             # Read source code path
             source_fp = pathlib.Path(args.source)
             # Read source code content
             source_code = read_from_path(fp=source_fp, is_json=False)
-            # API request timeout
-            timeout = 800
 
             # --- Panel Initialization ---
             summary_panel = SummaryPanel(
@@ -310,9 +307,8 @@ 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.<extension>)
-                runs_copy_source_fp = runs_dir / f"original{source_fp.suffix}"  # Use correct suffix
-                write_to_path(fp=runs_copy_source_fp, content=source_code)
+                # 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:
@@ -380,7 +376,8 @@ def main() -> None:
                     transition_delay=0.1,
                 )
 
-                for step in range(1, steps):
+                # Starting from step 1 to steps (inclusive) because the baseline solution is step 0, so we want to optimize for steps worth of steps
+                for step in range(1, steps + 1):
                     # 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
@@ -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"])
@@ -553,9 +544,7 @@ def main() -> None:
                     best_solution_score = None
 
                 if best_solution_code is None or best_solution_score is None:
-                    best_solution_content = (
-                        f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_copy_source_fp, is_json=False)}"
-                    )
+                    best_solution_content = f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_dir / f'step_0{source_fp.suffix}', is_json=False)}"
                 else:
                     # Format score for the comment
                     best_score_str = (
@@ -577,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.14.dist-info → weco-0.2.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.14
+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.14.dist-info → weco-0.2.16.dist-info}/WHEEL

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

weco-0.2.14.dist-info/RECORD

@@ -1,12 +0,0 @@
-weco/__init__.py,sha256=q4zeQ8CJq8NJyRcBMmST6zCzmK_HjNSFChrNBzJ9oks,426
-weco/api.py,sha256=z2DCe0kQJaaBHo-Vml52GRc9nwQBnrKATXoY3UfPojw,3842
-weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
-weco/cli.py,sha256=7pGJiyoBO7n6Xocwusx9iVqsaPj8OuPaJ1A2ks83Ekw,28961
-weco/panels.py,sha256=gB4rZbCvqzewUCBcILvyyU4fnOQLwFgHCGmtn-ZlgSo,13385
-weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
-weco-0.2.14.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.14.dist-info/METADATA,sha256=855avaVBC-wUyFDWIvITw0t4o_v4s90Us-t_3lQnPGw,10851
-weco-0.2.14.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-weco-0.2.14.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.14.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.14.dist-info/RECORD,,