feat: Implement User Authentication and Session Ownership

.github/workflows/release.yml

@@ -5,6 +5,8 @@ on:
     workflows: ["Lint and Format Code"]
     types:
       - completed
+    branches:
+      - main
 
   release:
     types: [published]

.gitignore

@@ -70,4 +70,7 @@ etc/
 digest.txt
 .runs/
 
-*.pyc
+*.pyc
+
+# Repomix
+repomix-output.*

.repomixignore

@@ -0,0 +1,4 @@
+examples/
+LICENSE
+.gitignore
+.repomixignore

README.md

@@ -10,7 +10,7 @@ Example applications include:
 
 - **GPU Kernel Optimization**: Reimplement PyTorch functions using CUDA, Triton or Metal, 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`
+- **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
 
@@ -20,7 +20,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)
 
 ---
 
@@ -32,25 +32,50 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
     pip install weco
     ```
 
-2.  **Configure API Keys:**
+2.  **Set Up LLM API Keys (Required):**
 
-    Set the appropriate environment variables for your desired language model provider:
+    `weco` requires API keys for the Large Language Models (LLMs) it uses internally. You **must** provide these keys via environment variables:
 
-    -   **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
-    -   **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.)
+    - **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
+    - **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.
+
 <div style="background-color: #fff3cd; border: 1px solid #ffeeba; padding: 15px; border-radius: 4px; margin-bottom: 15px;">
   <strong>⚠️ Warning: Code Modification</strong><br>
   <code>weco</code> directly modifies the file specified by <code>--source</code> during the optimization process. It is <strong>strongly recommended</strong> to use version control (like Git) to track changes and revert if needed. Alternatively, ensure you have a backup of your original file before running the command. Upon completion, the file will contain the best-performing version of the code found during the run.
 </div>
 
 ---
 
-### Example: Optimizing Simple PyTorch Operations
+### `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.
 
@@ -64,7 +89,7 @@ cd examples/hello-kernel-world
 pip install torch
 
 # Run Weco
-weco --source optimize.py \
+weco run --source optimize.py \
      --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
      --metric speedup \
      --maximize true \
@@ -77,19 +102,29 @@ weco --source optimize.py \
 
 ---
 
-### Command Line Arguments
+**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.7-sonnet`). Recommended models to try include `o4-mini`, 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       |
+| 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       |
+| `--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 logout` Command
+
+This command logs you out by removing the locally stored Weco API key.
+
+```bash
+weco logout
+```
 
 ---
 
@@ -98,6 +133,7 @@ weco --source optimize.py \
 Weco, powered by the AIDE algorithm, optimizes code iteratively based on your evaluation results. Achieving significant improvements, especially on complex research-level tasks, often requires substantial exploration time.
 
 The following plot from the independent [Research Engineering Benchmark (RE-Bench)](https://metr.org/AI_R_D_Evaluation_Report.pdf) report shows the performance of AIDE (the algorithm behind Weco) on challenging ML research engineering tasks over different time budgets.
+
 <p align="center">
 <img src="https://github.com/user-attachments/assets/ff0e471d-2f50-4e2d-b718-874862f533df" alt="RE-Bench Performance Across Time" width="60%"/>
 </p>
@@ -124,23 +160,25 @@ 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
 
 We welcome contributions! To get started:
 
 1.  **Fork and Clone the Repository:**
+
     ```bash
     git clone https://github.com/WecoAI/weco-cli.git
     cd weco-cli
     ```
 
 2.  **Install Development Dependencies:**
+
     ```bash
     pip install -e ".[dev]"
     ```
 
 3.  **Create a Feature Branch:**
+
     ```bash
     git checkout -b feature/your-feature-name
     ```

weco/__init__.py

@@ -1,4 +1,14 @@
+import os
+
 # DO NOT EDIT
 __pkg_version__ = "0.2.12"
 __api_version__ = "v1"
+
 __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"
+# If user specifies a custom base URL, use that instead
+if os.environ.get("WECO_BASE_URL"):
+    __base_url__ = os.environ.get("WECO_BASE_URL")
+
+__dashboard_url__ = "https://dashboard.weco.ai"
+if os.environ.get("WECO_DASHBOARD_URL"):
+    __dashboard_url__ = os.environ.get("WECO_DASHBOARD_URL")

weco/api.py

@@ -24,13 +24,15 @@ def start_optimization_session(
     search_policy_config: Dict[str, Any],
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
+    auth_headers: dict = {},  # Add auth_headers
     timeout: int = 800,
 ) -> 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",
+                f"{__base_url__}/sessions",  # Path is relative to base_url
                 json={
                     "source_code": source_code,
                     "additional_instructions": additional_instructions,
@@ -43,6 +45,7 @@ def start_optimization_session(
                     },
                     "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
                 },
+                headers=auth_headers,  # Add headers
                 timeout=timeout,
             )
             response.raise_for_status()
@@ -57,17 +60,20 @@ def evaluate_feedback_then_suggest_next_solution(
     execution_output: str,
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
+    auth_headers: dict = {},  # Add auth_headers
     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",
+            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()
@@ -77,12 +83,20 @@ def evaluate_feedback_then_suggest_next_solution(
 
 
 def get_optimization_session_status(
-    console: rich.console.Console, session_id: str, include_history: bool = False, timeout: int = 800
+    console: rich.console.Console,
+    session_id: str,
+    include_history: bool = False,
+    auth_headers: dict = {},
+    timeout: int = 800,  # Add auth_headers
 ) -> 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}", params={"include_history": include_history}, timeout=timeout
+            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()

weco/auth.py

@@ -0,0 +1,64 @@
+# weco/auth.py
+import os
+import pathlib
+import json
+import stat
+
+CONFIG_DIR = pathlib.Path.home() / ".config" / "weco"
+CREDENTIALS_FILE = CONFIG_DIR / "credentials.json"
+
+
+def ensure_config_dir():
+    """Ensures the configuration directory exists."""
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    # Ensure directory permissions are secure (optional but good practice)
+    try:
+        os.chmod(CONFIG_DIR, stat.S_IRWXU)  # Read/Write/Execute for owner only
+    except OSError as e:
+        print(f"Warning: Could not set permissions on {CONFIG_DIR}: {e}")
+
+
+def save_api_key(api_key: str):
+    """Saves the Weco API key securely."""
+    ensure_config_dir()
+    credentials = {"api_key": api_key}
+    try:
+        with open(CREDENTIALS_FILE, "w") as f:
+            json.dump(credentials, f)
+        # Set file permissions to read/write for owner only (600)
+        os.chmod(CREDENTIALS_FILE, stat.S_IRUSR | stat.S_IWUSR)
+    except IOError as e:
+        print(f"Error: Could not write credentials file at {CREDENTIALS_FILE}: {e}")
+    except OSError as e:
+        print(f"Warning: Could not set permissions on {CREDENTIALS_FILE}: {e}")
+
+
+def load_weco_api_key() -> str | None:
+    """Loads the Weco API key."""
+    if not CREDENTIALS_FILE.exists():
+        return None
+    try:
+        # Check permissions before reading (optional but safer)
+        file_stat = os.stat(CREDENTIALS_FILE)
+        if file_stat.st_mode & (stat.S_IRWXG | stat.S_IRWXO):  # Check if group/other have permissions
+            print(f"Warning: Credentials file {CREDENTIALS_FILE} has insecure permissions. Please set to 600.")
+            # Optionally, refuse to load or try to fix permissions
+
+        with open(CREDENTIALS_FILE, "r") as f:
+            credentials = json.load(f)
+            return credentials.get("api_key")
+    except (IOError, json.JSONDecodeError, OSError) as e:
+        print(f"Warning: Could not read or parse credentials file at {CREDENTIALS_FILE}: {e}")
+        return None
+
+
+def clear_api_key():
+    """Removes the stored API key."""
+    if CREDENTIALS_FILE.exists():
+        try:
+            os.remove(CREDENTIALS_FILE)
+            print("Logged out successfully.")
+        except OSError as e:
+            print(f"Error: Could not remove credentials file at {CREDENTIALS_FILE}: {e}")
+    else:
+        print("Already logged out.")