feat: Implement User Authentication and Session Ownership

README.md

@@ -10,9 +10,9 @@ 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
+![Weco CLI Demo](https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81)
 
 ---
 
@@ -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,10 @@
+import os
+
 # DO NOT EDIT
 __pkg_version__ = "0.2.8"
 __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")

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,
+                "additional_instructions": additional_instructions,  # Restore this field
                 "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_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.")