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

weco/__init__.py

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

weco/api.py

@@ -1,14 +1,20 @@
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 import rich
 import requests
 from weco import __pkg_version__, __base_url__
 import sys
+from rich.console import Console
 
 
 def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
     """Extract and display error messages from API responses in a structured format."""
-    console.print(f"[bold red]{e.response.json()['detail']}[/]")
-    sys.exit(1)
+    try:
+        detail = e.response.json()["detail"]
+    except (ValueError, KeyError):  # Handle cases where response is not JSON or detail key is missing
+        detail = f"HTTP {e.response.status_code} Error: {e.response.text}"
+    console.print(f"[bold red]{detail}[/]")
+    # Avoid exiting here, let the caller decide if the error is fatal
+    # sys.exit(1)
 
 
 def start_optimization_session(
@@ -28,25 +34,32 @@ def start_optimization_session(
 ) -> Dict[str, Any]:
     """Start the optimization session."""
     with console.status("[bold green]Starting Optimization..."):
-        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,
+        try:
+            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},
                 },
-                "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
-            },
-            headers=auth_headers,  # Add headers
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
+                headers=auth_headers,  # Add headers
+                timeout=timeout,
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.HTTPError as e:
+            handle_api_error(e, console)
+            sys.exit(1)  # Exit if starting session fails
+        except requests.exceptions.RequestException as e:
+            console.print(f"[bold red]Network Error starting session: {e}[/]")
+            sys.exit(1)
 
 
 def evaluate_feedback_then_suggest_next_solution(
@@ -58,29 +71,92 @@ def evaluate_feedback_then_suggest_next_solution(
     timeout: int = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
-    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()
+    try:
+        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:
+        # Allow caller to handle suggest errors, maybe retry or terminate
+        handle_api_error(e, Console())  # Use default console if none passed
+        raise  # Re-raise the exception
+    except requests.exceptions.RequestException as e:
+        print(f"Network Error during suggest: {e}")  # Use print as console might not be available
+        raise  # Re-raise the exception
 
 
 def get_optimization_session_status(
     session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
 ) -> Dict[str, Any]:
     """Get the current status of the optimization session."""
-    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()
+    try:
+        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()
+    except requests.exceptions.HTTPError as e:
+        handle_api_error(e, Console())  # Use default console
+        raise  # Re-raise
+    except requests.exceptions.RequestException as e:
+        print(f"Network Error getting status: {e}")
+        raise  # Re-raise
+
+
+def send_heartbeat(
+    session_id: str,
+    auth_headers: dict = {},
+    timeout: int = 10,  # Shorter timeout for non-critical heartbeat
+) -> bool:
+    """Send a heartbeat signal to the backend."""
+    try:
+        response = requests.put(f"{__base_url__}/sessions/{session_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
+        return True
+    except requests.exceptions.HTTPError as e:
+        # Log non-critical errors like 409 Conflict (session not running)
+        if e.response.status_code == 409:
+            print(f"Heartbeat ignored: Session {session_id} is not running.", file=sys.stderr)
+        else:
+            print(f"Heartbeat failed for session {session_id}: HTTP {e.response.status_code}", file=sys.stderr)
+        # Don't exit, just report failure
+        return False
+    except requests.exceptions.RequestException as e:
+        # Network errors are also non-fatal for heartbeats
+        print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
+        return False
+
+
+def report_termination(
+    session_id: str,
+    status_update: str,
+    reason: str,
+    details: Optional[str] = None,
+    auth_headers: dict = {},
+    timeout: int = 30,  # Reasonably longer timeout for important termination message
+) -> bool:
+    """Report the termination reason to the backend."""
+    try:
+        response = requests.post(
+            f"{__base_url__}/sessions/{session_id}/terminate",
+            json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
+            headers=auth_headers,
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        return True
+    except requests.exceptions.RequestException as e:
+        # Log failure, but don't prevent CLI exit
+        print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
+        return False

weco/cli.py

@@ -5,6 +5,9 @@ import math
 import time
 import requests
 import webbrowser
+import threading
+import signal
+import traceback
 from rich.console import Console
 from rich.live import Live
 from rich.panel import Panel
@@ -15,6 +18,8 @@ from .api import (
     evaluate_feedback_then_suggest_next_solution,
     get_optimization_session_status,
     handle_api_error,
+    send_heartbeat,
+    report_termination,
 )
 
 from . import __base_url__
@@ -37,11 +42,73 @@ from .utils import (
     run_evaluation,
     smooth_update,
     format_number,
+    check_for_cli_updates,
 )
 
 install(show_locals=True)
 console = Console()
 
+# --- Global variable for heartbeat thread ---
+heartbeat_thread = None
+stop_heartbeat_event = threading.Event()
+current_session_id_for_heartbeat = None
+current_auth_headers_for_heartbeat = {}
+
+
+# --- Heartbeat Sender Class ---
+class HeartbeatSender(threading.Thread):
+    def __init__(self, session_id: str, auth_headers: dict, stop_event: threading.Event, interval: int = 30):
+        super().__init__(daemon=True)  # Daemon thread exits when main thread exits
+        self.session_id = session_id
+        self.auth_headers = auth_headers
+        self.interval = interval
+        self.stop_event = stop_event
+
+    def run(self):
+        try:
+            while not self.stop_event.is_set():
+                if not send_heartbeat(self.session_id, self.auth_headers):
+                    # send_heartbeat itself prints errors to stderr if it returns False
+                    # No explicit HeartbeatSender log needed here unless more detail is desired for a False return
+                    pass  # Continue trying as per original logic
+
+                if self.stop_event.is_set():  # Check before waiting for responsiveness
+                    break
+
+                self.stop_event.wait(self.interval)  # Wait for interval or stop signal
+
+        except Exception as e:
+            # Catch any unexpected error in the loop to prevent silent thread death
+            print(
+                f"[ERROR HeartbeatSender] Unhandled exception in run loop for session {self.session_id}: {e}", file=sys.stderr
+            )
+            traceback.print_exc(file=sys.stderr)
+            # The loop will break due to the exception, and thread will terminate via finally.
+
+
+# --- Signal Handling ---
+def signal_handler(signum, frame):
+    signal_name = signal.Signals(signum).name
+    console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]")
+
+    # Stop heartbeat thread
+    stop_heartbeat_event.set()
+    if heartbeat_thread and heartbeat_thread.is_alive():
+        heartbeat_thread.join(timeout=2)  # Give it a moment to stop
+
+    # Report termination (best effort)
+    if current_session_id_for_heartbeat:
+        report_termination(
+            session_id=current_session_id_for_heartbeat,
+            status_update="terminated",
+            reason=f"user_terminated_{signal_name.lower()}",
+            details=f"Process terminated by signal {signal_name} ({signum}).",
+            auth_headers=current_auth_headers_for_heartbeat,
+        )
+
+    # Exit gracefully
+    sys.exit(0)
+
 
 def perform_login(console: Console):
     """Handles the device login flow."""
@@ -161,6 +228,16 @@ def perform_login(console: Console):
 
 def main() -> None:
     """Main function for the Weco CLI."""
+    # Setup signal handlers
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    # --- Perform Update Check ---
+    # Import __pkg_version__ here to avoid circular import issues if it's also used in modules imported by cli.py
+    from . import __pkg_version__
+
+    check_for_cli_updates(__pkg_version__)  # Call the imported function
+
     # --- Argument Parsing ---
     parser = argparse.ArgumentParser(
         description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
@@ -170,29 +247,55 @@ def main() -> None:
 
     # --- Run Command ---
     run_parser = subparsers.add_parser(
-        "run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter
+        "run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False
     )
     # Add arguments specific to the 'run' command to the run_parser
-    run_parser.add_argument("--source", type=str, required=True, help="Path to the source code file (e.g. optimize.py)")
     run_parser.add_argument(
-        "--eval-command", type=str, required=True, help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1')"
+        "-s",
+        "--source",
+        type=str,
+        required=True,
+        help="Path to the source code file that will be optimized (e.g., `optimize.py`)",
     )
-    run_parser.add_argument("--metric", type=str, required=True, help="Metric to optimize")
     run_parser.add_argument(
-        "--maximize",
+        "-c",
+        "--eval-command",
         type=str,
-        choices=["true", "false"],
         required=True,
-        help="Specify 'true' to maximize the metric or 'false' to minimize.",
+        help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1').",
     )
-    run_parser.add_argument("--steps", type=int, required=True, help="Number of steps to run")
-    run_parser.add_argument("--model", type=str, required=True, help="Model to use for optimization")
-    run_parser.add_argument("--log-dir", type=str, default=".runs", help="Directory to store logs and results")
     run_parser.add_argument(
+        "-m",
+        "--metric",
+        type=str,
+        required=True,
+        help="Metric to optimize (e.g. 'accuracy', 'loss', 'f1_score') that is printed to the terminal by the eval command.",
+    )
+    run_parser.add_argument(
+        "-g",
+        "--goal",
+        type=str,
+        choices=["maximize", "max", "minimize", "min"],
+        required=True,
+        help="Specify 'maximize'/'max' to maximize the metric or 'minimize'/'min' to minimize it.",
+    )
+    run_parser.add_argument("-n", "--steps", type=int, default=100, help="Number of steps to run. Defaults to 100.")
+    run_parser.add_argument(
+        "-M",
+        "--model",
+        type=str,
+        default=None,
+        help="Model to use for optimization. Defaults to `o4-mini` when `OPENAI_API_KEY` is set, `claude-3-7-sonnet-20250219` when `ANTHROPIC_API_KEY` is set, and `gemini-2.5-pro-exp-03-25` when `GEMINI_API_KEY` is set. When multiple keys are set, the priority is `OPENAI_API_KEY` > `ANTHROPIC_API_KEY` > `GEMINI_API_KEY`.",
+    )
+    run_parser.add_argument(
+        "-l", "--log-dir", type=str, default=".runs", help="Directory to store logs and results. Defaults to `.runs`."
+    )
+    run_parser.add_argument(
+        "-i",
         "--additional-instructions",
         default=None,
         type=str,
-        help="Description of additional instruction or path to a file containing additional instructions",
+        help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
     )
 
     # --- Logout Command ---
@@ -207,6 +310,10 @@ def main() -> None:
 
     # --- Handle Run Command ---
     elif args.command == "run":
+        global heartbeat_thread, current_session_id_for_heartbeat, current_auth_headers_for_heartbeat  # Allow modification of globals
+
+        session_id = None  # Initialize session_id
+        optimization_completed_normally = False  # Flag for finally block
         # --- Check Authentication ---
         weco_api_key = load_weco_api_key()
         llm_api_keys = read_api_keys_from_env()  # Read keys from client environment
@@ -238,18 +345,29 @@ def main() -> None:
 
         # --- Prepare API Call Arguments ---
         auth_headers = {}
-
         if weco_api_key:
             auth_headers["Authorization"] = f"Bearer {weco_api_key}"
-            # Backend will decide whether to use client keys based on auth status
+        current_auth_headers_for_heartbeat = auth_headers  # Store for signal handler
 
         # --- Main Run Logic ---
         try:
             # --- Configuration Loading ---
             evaluation_command = args.eval_command
             metric_name = args.metric
-            maximize = args.maximize == "true"
+            maximize = args.goal in ["maximize", "max"]
             steps = args.steps
+            # Determine the model to use
+            if args.model is None:
+                if "OPENAI_API_KEY" in llm_api_keys:
+                    args.model = "o4-mini"
+                elif "ANTHROPIC_API_KEY" in llm_api_keys:
+                    args.model = "claude-3-7-sonnet-20250219"
+                elif "GEMINI_API_KEY" in llm_api_keys:
+                    args.model = "gemini-2.5-pro-exp-03-25"
+                else:
+                    raise ValueError(
+                        "No LLM API keys found in environment. Please set one of the following: OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY."
+                    )
             code_generator_config = {"model": args.model}
             evaluator_config = {"model": args.model, "include_analysis": True}
             search_policy_config = {
@@ -289,16 +407,22 @@ def main() -> None:
                 evaluator_config=evaluator_config,
                 search_policy_config=search_policy_config,
                 additional_instructions=additional_instructions,
-                api_keys=llm_api_keys,  # Pass client LLM keys
-                auth_headers=auth_headers,  # Pass Weco key if logged in
+                api_keys=llm_api_keys,
+                auth_headers=auth_headers,
                 timeout=timeout,
             )
+            session_id = session_response["session_id"]
+            current_session_id_for_heartbeat = session_id  # Store for signal handler/finally
+
+            # --- Start Heartbeat Thread ---
+            stop_heartbeat_event.clear()  # Ensure event is clear before starting
+            heartbeat_thread = HeartbeatSender(session_id, auth_headers, stop_heartbeat_event)
+            heartbeat_thread.start()
 
             # --- Live Update Loop ---
             refresh_rate = 4
             with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
                 # Define the runs directory (.runs/<session-id>)
-                session_id = session_response["session_id"]
                 runs_dir = pathlib.Path(args.log_dir) / session_id
                 runs_dir.mkdir(parents=True, exist_ok=True)
 
@@ -358,6 +482,9 @@ def main() -> None:
                     transition_delay=0.1,
                 )
 
+                # # Send initial heartbeat immediately after starting
+                # send_heartbeat(session_id, auth_headers)
+
                 # Run evaluation on the initial solution
                 term_out = run_evaluation(eval_command=args.eval_command)
 
@@ -386,6 +513,7 @@ def main() -> None:
                         auth_headers=auth_headers,  # Pass Weco key if logged in
                         timeout=timeout,
                     )
+
                     # 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"]
@@ -476,7 +604,7 @@ def main() -> None:
                     additional_instructions=args.additional_instructions
                 )
 
-                # Ensure we pass evaluation results for the last step's generated solution
+                # Final evaluation report
                 eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
                     session_id=session_id,
                     execution_output=term_out,
@@ -555,14 +683,74 @@ def main() -> None:
                 # write the best solution to the source file
                 write_to_path(fp=source_fp, content=best_solution_content)
 
+                # Mark as completed normally for the finally block
+                optimization_completed_normally = True
+
             console.print(end_optimization_layout)
 
         except Exception as e:
+            # Catch errors during the main optimization loop or setup
             try:
-                error_message = e.response.json()["detail"]
+                error_message = e.response.json()["detail"]  # Try to get API error 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
+                error_message = str(e)  # Otherwise, use the exception string
+            console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Optimization Error", border_style="red"))
+            # Print traceback for debugging if needed (can be noisy)
             # console.print_exception(show_locals=False)
-            sys.exit(1)
+
+            # Ensure optimization_completed_normally is False
+            optimization_completed_normally = False
+
+            # Prepare details for termination report
+            error_details = traceback.format_exc()
+
+            # Exit code will be handled by finally block or sys.exit below
+            exit_code = 1  # Indicate error
+            # No sys.exit here, let finally block run
+
+        finally:
+            # This block runs whether the try block completed normally or raised an exception
+
+            # Stop heartbeat thread
+            stop_heartbeat_event.set()
+            if heartbeat_thread and heartbeat_thread.is_alive():
+                heartbeat_thread.join(timeout=2)  # Give it a moment to stop
+
+            # Report final status if a session was started
+            if session_id:
+                final_status = "unknown"
+                final_reason = "unknown_termination"
+                final_details = None
+
+                if optimization_completed_normally:
+                    final_status = "completed"
+                    final_reason = "completed_successfully"
+                else:
+                    # If an exception was caught and we have details
+                    if "error_details" in locals():
+                        final_status = "error"
+                        final_reason = "error_cli_internal"
+                        final_details = error_details
+                    # else: # Should have been handled by signal handler if terminated by user
+                    # Keep default 'unknown' if we somehow end up here without error/completion/signal
+
+                # Avoid reporting if terminated by signal handler (already reported)
+                # Check a flag or rely on status not being 'unknown'
+                if final_status != "unknown":
+                    report_termination(
+                        session_id=session_id,
+                        status_update=final_status,
+                        reason=final_reason,
+                        details=final_details,
+                        auth_headers=auth_headers,
+                    )
+
+            # Ensure proper exit code if an error occurred
+            if not optimization_completed_normally and "exit_code" in locals() and exit_code != 0:
+                sys.exit(exit_code)
+            elif not optimization_completed_normally:
+                # Generic error exit if no specific code was set but try block failed
+                sys.exit(1)
+            else:
+                # Normal exit
+                sys.exit(0)

weco/panels.py

@@ -253,7 +253,7 @@ class MetricTreePanel:
         # Make sure the metric tree is built before calling build_rich_tree
         return Panel(
             self._build_rich_tree(),
-            title="[bold]🔎 Exploring Solutions..." if not is_done else "[bold]🔎 Optimization Complete!",
+            title=("[bold]🔎 Exploring Solutions..." if not is_done else "[bold]🔎 Optimization Complete!"),
             border_style="green",
             expand=True,
             padding=(0, 1),

weco/utils.py

@@ -7,6 +7,8 @@ from rich.layout import Layout
 from rich.live import Live
 from rich.panel import Panel
 import pathlib
+import requests
+from packaging.version import parse as parse_version
 
 
 # Env/arg helper functions
@@ -109,3 +111,33 @@ def run_evaluation(eval_command: str) -> str:
             output += "\n"
         output += result.stdout
     return output
+
+
+# Update Check Function
+def check_for_cli_updates(current_version_str: str):
+    """Checks PyPI for a newer version of the weco package and notifies the user."""
+    try:
+        pypi_url = "https://pypi.org/pypi/weco/json"
+        response = requests.get(pypi_url, timeout=5)  # Short timeout for non-critical check
+        response.raise_for_status()
+        latest_version_str = response.json()["info"]["version"]
+
+        current_version = parse_version(current_version_str)
+        latest_version = parse_version(latest_version_str)
+
+        if latest_version > current_version:
+            yellow_start = "\033[93m"
+            reset_color = "\033[0m"
+            message = f"WARNING: New weco version ({latest_version_str}) available (you have {current_version_str}). Run: pip install --upgrade weco"
+            print(f"{yellow_start}{message}{reset_color}")
+            time.sleep(2)  # Wait for 2 second
+
+    except requests.exceptions.RequestException:
+        # Silently fail on network errors, etc. Don't disrupt user.
+        pass
+    except (KeyError, ValueError):
+        # Handle cases where the PyPI response format might be unexpected
+        pass
+    except Exception:
+        # Catch any other unexpected error during the check
+        pass

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.17
+Version: 0.2.19
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -14,6 +14,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: rich
+Requires-Dist: packaging
 Provides-Extra: dev
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: build; extra == "dev"
@@ -22,15 +23,13 @@ Dynamic: license-file
 
 <div align="center">
 
-# Weco: The AI Code Optimizer
+# Weco: The Platform for Self-Improving Code
 
 [![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>
 
 ---
@@ -39,9 +38,9 @@ Weco systematically optimizes your code, guided directly by your evaluation metr
 
 Example applications include:
 
-- **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`
+- **GPU Kernel Optimization**: Reimplement PyTorch functions using [CUDA](/examples/cuda/README.md) or [Triton](/examples/triton/README.md), optimizing for `latency`, `throughput`, or `memory_bandwidth`.
+- **Model Development**: Tune feature transformations, architectures or [the whole training pipeline](/examples/spaceship-titanic/README.md), optimizing for `validation_accuracy`, `AUC`, or `Sharpe Ratio`.
+- **Prompt Engineering**: Refine prompts for LLMs (e.g., for [math problems](/examples/prompt/README.md)), optimizing for `win_rate`, `relevance`, or `format_adherence`
 
 ![image](assets/example-optimization.gif)
 
@@ -71,29 +70,9 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
     - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
     - **Google DeepMind:** `export GEMINI_API_KEY="your_key_here"` (Google AI Studio has a free API usage quota. Create a key [here](https://aistudio.google.com/apikey) to use `weco` for free.)
 
-    The optimization process will fail if the necessary keys for the chosen model are not found in your environment.
-
-3.  **Log In to Weco (Optional):**
-
-    To associate your optimization runs with your Weco account and view them on the Weco dashboard, you can log in. `weco` uses a device authentication flow:
-
-    - When you first run `weco run`, you'll be prompted if you want to log in or proceed anonymously.
-    - If you choose to log in (by pressing `l`), you'll be shown a URL and `weco` will attempt to open it in your default web browser.
-    - You then authenticate in the browser. Once authenticated, the CLI will detect this and complete the login.
-    - This saves a Weco-specific API key locally (typically at `~/.config/weco/credentials.json`).
-
-    If you choose to skip login (by pressing Enter or `s`), `weco` will still function using the environment variable LLM keys, but the run history will not be linked to a Weco account.
-
-    To log out and remove your saved Weco API key, use the `weco logout` command.
-
 ---
 
-## Usage
-
-The CLI has two main commands:
-
-- `weco run`: Initiates the code optimization process.
-- `weco logout`: Logs you out of your Weco account.
+## Get Started
 
 <div style="background-color: #fff3cd; border: 1px solid #ffeeba; padding: 15px; border-radius: 4px; margin-bottom: 15px;">
   <strong>⚠️ Warning: Code Modification</strong><br>
@@ -102,10 +81,6 @@ The CLI has two main commands:
 
 ---
 
-### `weco run` Command
-
-This command starts the optimization process.
-
 **Example: Optimizing Simple PyTorch Operations**
 
 This basic example shows how to optimize a simple PyTorch function for speedup.
@@ -123,9 +98,8 @@ pip install torch
 weco run --source optimize.py \
      --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
      --metric speedup \
-     --maximize true \
+     --goal maximize \
      --steps 15 \
-     --model gemini-2.5-pro-exp-03-25 \
      --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
 ```
 
@@ -133,28 +107,33 @@ weco run --source optimize.py \
 
 ---
 
-**Arguments for `weco run`:**
+### Arguments for `weco run`
 
-| Argument                    | Description                                                                                                                                                               | Required |
-| :-------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | :------- |
-| `--source`                  | Path to the source code file that will be optimized (e.g., `optimize.py`).                                                                                                | Yes      |
-| `--eval-command`            | Command to run for evaluating the code in `--source`. This command should print the target `--metric` and its value to the terminal (stdout/stderr). See note below.      | Yes      |
-| `--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      |
-| `--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       |
+**Required:**
 
----
+| Argument            | Description                                                                                                                                                                                  |
+| :------------------ | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `-s, --source`      | Path to the source code file that will be optimized (e.g., `optimize.py`).                                                                                                                   |
+| `-c, --eval-command`| Command to run for evaluating the code in `--source`. This command should print the target `--metric` and its value to the terminal (stdout/stderr). See note below.                        |
+| `-m, --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`.                                    |
+| `-g, --goal`   | `maximize`/`max` to maximize the `--metric` or `minimize`/`min` to minimize it. |
 
-### `weco logout` Command
+<br>
 
-This command logs you out by removing the locally stored Weco API key.
+**Optional:**
 
-```bash
-weco logout
-```
+| Argument                       | Description                                                                                                                                                                                                                | Default |
+| :----------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :------ |
+| `-n, --steps`                  | Number of optimization steps (LLM iterations) to run.                                                                                                                                                                      | 100 |
+| `-M, --model`                  | Model identifier for the LLM to use (e.g., `gpt-4o`, `claude-3.5-sonnet`).                                                                                                                                                 | `o4-mini` when `OPENAI_API_KEY` is set; `claude-3-7-sonnet-20250219` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro-exp-03-25` when `GEMINI_API_KEY` is set (priority: `OPENAI_API_KEY` > `ANTHROPIC_API_KEY` > `GEMINI_API_KEY`). |
+| `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM.                                                                                             | `None` |
+| `-l, --log-dir`                | Path to the directory to log intermediate steps and final optimization result.                                                                                                                                             | `.runs/` |
+
+---
+
+### Weco Dashboard
+To associate your optimization runs with your Weco account and view them on the Weco dashboard, you can log in. `weco` uses a device authentication flow
+![image (16)](https://github.com/user-attachments/assets/8a0a285b-4894-46fa-b6a2-4990017ca0c6)
 
 ---
 

weco-0.2.19.dist-info/RECORD

@@ -0,0 +1,12 @@
+weco/__init__.py,sha256=npWmRgLxfVK69GdyxIujnI87xqmPCBrZWxxAxL_QQOc,478
+weco/api.py,sha256=lJJ0j0-bABiQXDlRb43fCo7ky0N_HwfZgFdMktRKQ90,6635
+weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
+weco/cli.py,sha256=eI468fxpMTfGPL-aX6EMYxh0NuaRxpaLVF_Jj2DiFhU,36383
+weco/panels.py,sha256=pM_YGnmcXM_1CBcxo_EAzOV3g_4NFdLS4MqDqx7THbA,13563
+weco/utils.py,sha256=LVTBo3dduJmhlbotcYoUW2nLx6IRtKs4eDFR52Qltcg,5244
+weco-0.2.19.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.19.dist-info/METADATA,sha256=3VBVsCqr7p332A10KsLr168GvOIKcCOWWfGDv8ViF7I,10729
+weco-0.2.19.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+weco-0.2.19.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.19.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.19.dist-info/RECORD,,

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

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

weco-0.2.17.dist-info/RECORD

@@ -1,12 +0,0 @@
-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,,