weco 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
weco/__init__.py CHANGED
@@ -1,4 +1,14 @@
1
+ import os
2
+
1
3
  # DO NOT EDIT
2
- __pkg_version__ = "0.2.12"
4
+ __pkg_version__ = "0.2.14"
3
5
  __api_version__ = "v1"
4
- __base_url__ = f"https://api.aide.weco.ai/{__api_version__}"
6
+
7
+ __base_url__ = f"https://api.weco.ai/{__api_version__}"
8
+ # If user specifies a custom base URL, use that instead
9
+ if os.environ.get("WECO_BASE_URL"):
10
+ __base_url__ = os.environ.get("WECO_BASE_URL")
11
+
12
+ __dashboard_url__ = "https://dashboard.weco.ai"
13
+ if os.environ.get("WECO_DASHBOARD_URL"):
14
+ __dashboard_url__ = os.environ.get("WECO_DASHBOARD_URL")
weco/api.py CHANGED
@@ -7,8 +7,7 @@ import sys
7
7
 
8
8
  def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
9
9
  """Extract and display error messages from API responses in a structured format."""
10
- error_message = str(e) # Default message
11
- console.print(f"[bold red]Error:[/] {error_message}")
10
+ console.print(f"[bold red]{e.response.json()['detail']}[/]")
12
11
  sys.exit(1)
13
12
 
14
13
 
@@ -24,13 +23,15 @@ def start_optimization_session(
24
23
  search_policy_config: Dict[str, Any],
25
24
  additional_instructions: str = None,
26
25
  api_keys: Dict[str, Any] = {},
26
+ auth_headers: dict = {}, # Add auth_headers
27
27
  timeout: int = 800,
28
28
  ) -> Dict[str, Any]:
29
29
  """Start the optimization session."""
30
30
  with console.status("[bold green]Starting Optimization..."):
31
31
  try:
32
+ # __base_url__ already contains /v1
32
33
  response = requests.post(
33
- f"{__base_url__}/sessions",
34
+ f"{__base_url__}/sessions", # Path is relative to base_url
34
35
  json={
35
36
  "source_code": source_code,
36
37
  "additional_instructions": additional_instructions,
@@ -43,6 +44,7 @@ def start_optimization_session(
43
44
  },
44
45
  "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
45
46
  },
47
+ headers=auth_headers, # Add headers
46
48
  timeout=timeout,
47
49
  )
48
50
  response.raise_for_status()
@@ -57,17 +59,20 @@ def evaluate_feedback_then_suggest_next_solution(
57
59
  execution_output: str,
58
60
  additional_instructions: str = None,
59
61
  api_keys: Dict[str, Any] = {},
62
+ auth_headers: dict = {}, # Add auth_headers
60
63
  timeout: int = 800,
61
64
  ) -> Dict[str, Any]:
62
65
  """Evaluate the feedback and suggest the next solution."""
63
66
  try:
67
+ # __base_url__ already contains /v1
64
68
  response = requests.post(
65
- f"{__base_url__}/sessions/{session_id}/suggest",
69
+ f"{__base_url__}/sessions/{session_id}/suggest", # Path is relative to base_url
66
70
  json={
67
71
  "execution_output": execution_output,
68
72
  "additional_instructions": additional_instructions,
69
73
  "metadata": {**api_keys},
70
74
  },
75
+ headers=auth_headers, # Add headers
71
76
  timeout=timeout,
72
77
  )
73
78
  response.raise_for_status()
@@ -77,12 +82,20 @@ def evaluate_feedback_then_suggest_next_solution(
77
82
 
78
83
 
79
84
  def get_optimization_session_status(
80
- console: rich.console.Console, session_id: str, include_history: bool = False, timeout: int = 800
85
+ console: rich.console.Console,
86
+ session_id: str,
87
+ include_history: bool = False,
88
+ auth_headers: dict = {},
89
+ timeout: int = 800, # Add auth_headers
81
90
  ) -> Dict[str, Any]:
82
91
  """Get the current status of the optimization session."""
83
92
  try:
93
+ # __base_url__ already contains /v1
84
94
  response = requests.get(
85
- f"{__base_url__}/sessions/{session_id}", params={"include_history": include_history}, timeout=timeout
95
+ f"{__base_url__}/sessions/{session_id}", # Path is relative to base_url
96
+ params={"include_history": include_history},
97
+ headers=auth_headers, # Add headers
98
+ timeout=timeout,
86
99
  )
87
100
  response.raise_for_status()
88
101
  return response.json()
weco/auth.py ADDED
@@ -0,0 +1,64 @@
1
+ # weco/auth.py
2
+ import os
3
+ import pathlib
4
+ import json
5
+ import stat
6
+
7
+ CONFIG_DIR = pathlib.Path.home() / ".config" / "weco"
8
+ CREDENTIALS_FILE = CONFIG_DIR / "credentials.json"
9
+
10
+
11
+ def ensure_config_dir():
12
+ """Ensures the configuration directory exists."""
13
+ CONFIG_DIR.mkdir(parents=True, exist_ok=True)
14
+ # Ensure directory permissions are secure (optional but good practice)
15
+ try:
16
+ os.chmod(CONFIG_DIR, stat.S_IRWXU) # Read/Write/Execute for owner only
17
+ except OSError as e:
18
+ print(f"Warning: Could not set permissions on {CONFIG_DIR}: {e}")
19
+
20
+
21
+ def save_api_key(api_key: str):
22
+ """Saves the Weco API key securely."""
23
+ ensure_config_dir()
24
+ credentials = {"api_key": api_key}
25
+ try:
26
+ with open(CREDENTIALS_FILE, "w") as f:
27
+ json.dump(credentials, f)
28
+ # Set file permissions to read/write for owner only (600)
29
+ os.chmod(CREDENTIALS_FILE, stat.S_IRUSR | stat.S_IWUSR)
30
+ except IOError as e:
31
+ print(f"Error: Could not write credentials file at {CREDENTIALS_FILE}: {e}")
32
+ except OSError as e:
33
+ print(f"Warning: Could not set permissions on {CREDENTIALS_FILE}: {e}")
34
+
35
+
36
+ def load_weco_api_key() -> str | None:
37
+ """Loads the Weco API key."""
38
+ if not CREDENTIALS_FILE.exists():
39
+ return None
40
+ try:
41
+ # Check permissions before reading (optional but safer)
42
+ file_stat = os.stat(CREDENTIALS_FILE)
43
+ if file_stat.st_mode & (stat.S_IRWXG | stat.S_IRWXO): # Check if group/other have permissions
44
+ print(f"Warning: Credentials file {CREDENTIALS_FILE} has insecure permissions. Please set to 600.")
45
+ # Optionally, refuse to load or try to fix permissions
46
+
47
+ with open(CREDENTIALS_FILE, "r") as f:
48
+ credentials = json.load(f)
49
+ return credentials.get("api_key")
50
+ except (IOError, json.JSONDecodeError, OSError) as e:
51
+ print(f"Warning: Could not read or parse credentials file at {CREDENTIALS_FILE}: {e}")
52
+ return None
53
+
54
+
55
+ def clear_api_key():
56
+ """Removes the stored API key."""
57
+ if CREDENTIALS_FILE.exists():
58
+ try:
59
+ os.remove(CREDENTIALS_FILE)
60
+ print("Logged out successfully.")
61
+ except OSError as e:
62
+ print(f"Error: Could not remove credentials file at {CREDENTIALS_FILE}: {e}")
63
+ else:
64
+ print("Already logged out.")
weco/cli.py CHANGED
@@ -2,11 +2,23 @@ import argparse
2
2
  import sys
3
3
  import pathlib
4
4
  import math
5
+ import time
6
+ import requests
7
+ import webbrowser
5
8
  from rich.console import Console
6
9
  from rich.live import Live
7
10
  from rich.panel import Panel
8
11
  from rich.traceback import install
9
- from .api import start_optimization_session, evaluate_feedback_then_suggest_next_solution, get_optimization_session_status
12
+ from rich.prompt import Prompt
13
+ from .api import (
14
+ start_optimization_session,
15
+ evaluate_feedback_then_suggest_next_solution,
16
+ get_optimization_session_status,
17
+ handle_api_error,
18
+ )
19
+
20
+ from . import __base_url__
21
+ from .auth import load_weco_api_key, save_api_key, clear_api_key
10
22
  from .panels import (
11
23
  SummaryPanel,
12
24
  PlanPanel,
@@ -31,348 +43,541 @@ install(show_locals=True)
31
43
  console = Console()
32
44
 
33
45
 
46
+ def perform_login(console: Console):
47
+ """Handles the device login flow."""
48
+ try:
49
+ # 1. Initiate device login
50
+ console.print("Initiating login...")
51
+ init_response = requests.post(f"{__base_url__}/auth/device/initiate")
52
+ init_response.raise_for_status()
53
+ init_data = init_response.json()
54
+
55
+ device_code = init_data["device_code"]
56
+ verification_uri = init_data["verification_uri"]
57
+ expires_in = init_data["expires_in"]
58
+ interval = init_data["interval"]
59
+
60
+ # 2. Display instructions
61
+ console.print("\n[bold yellow]Action Required:[/]")
62
+ console.print("Please open the following URL in your browser to authenticate:")
63
+ console.print(f"[link={verification_uri}]{verification_uri}[/link]")
64
+ console.print(f"This request will expire in {expires_in // 60} minutes.")
65
+ console.print("Attempting to open the authentication page in your default browser...") # Notify user
66
+
67
+ # Automatically open the browser
68
+ try:
69
+ if not webbrowser.open(verification_uri):
70
+ console.print("[yellow]Could not automatically open the browser. Please open the link manually.[/]")
71
+ except Exception as browser_err:
72
+ console.print(
73
+ f"[yellow]Could not automatically open the browser ({browser_err}). Please open the link manually.[/]"
74
+ )
75
+
76
+ console.print("Waiting for authentication...", end="")
77
+
78
+ # 3. Poll for token
79
+ start_time = time.time()
80
+ # Use a simple text update instead of Spinner within Live for potentially better compatibility
81
+ polling_status = "Waiting..."
82
+ with Live(polling_status, refresh_per_second=1, transient=True, console=console) as live_status:
83
+ while True:
84
+ # Check for timeout
85
+ if time.time() - start_time > expires_in:
86
+ console.print("\n[bold red]Error:[/] Login request timed out.")
87
+ return False
88
+
89
+ time.sleep(interval)
90
+ live_status.update("Waiting... (checking status)")
91
+
92
+ try:
93
+ token_response = requests.post(
94
+ f"{__base_url__}/auth/device/token", # REMOVED /v1 prefix
95
+ json={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code},
96
+ )
97
+
98
+ # Check for 202 Accepted - Authorization Pending
99
+ if token_response.status_code == 202:
100
+ token_data = token_response.json()
101
+ if token_data.get("error") == "authorization_pending":
102
+ live_status.update("Waiting... (authorization pending)")
103
+ continue # Continue polling
104
+ else:
105
+ # Unexpected 202 response format
106
+ console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
107
+ return False
108
+
109
+ # Check for standard OAuth2 errors (often 400 Bad Request)
110
+ elif token_response.status_code == 400:
111
+ token_data = token_response.json()
112
+ error_code = token_data.get("error", "unknown_error")
113
+ # NOTE: Removed "authorization_pending" check from here
114
+ if error_code == "slow_down":
115
+ interval += 5 # Increase polling interval if instructed
116
+ live_status.update(f"Waiting... (slowing down polling to {interval}s)")
117
+ continue
118
+ elif error_code == "expired_token":
119
+ console.print("\n[bold red]Error:[/] Login request expired.")
120
+ return False
121
+ elif error_code == "access_denied":
122
+ console.print("\n[bold red]Error:[/] Authorization denied by user.")
123
+ return False
124
+ else: # invalid_grant, etc.
125
+ error_desc = token_data.get("error_description", "Unknown error during polling.")
126
+ console.print(f"\n[bold red]Error:[/] {error_desc} ({error_code})")
127
+ return False
128
+
129
+ # Check for other non-200/non-202/non-400 HTTP errors
130
+ token_response.raise_for_status()
131
+
132
+ # If successful (200 OK and no 'error' field)
133
+ token_data = token_response.json()
134
+ if "access_token" in token_data:
135
+ api_key = token_data["access_token"]
136
+ save_api_key(api_key)
137
+ console.print("\n[bold green]Login successful![/]")
138
+ return True
139
+ else:
140
+ # Unexpected successful response format
141
+ console.print("\n[bold red]Error:[/] Received unexpected response from server during polling.")
142
+ print(token_data) # Log for debugging
143
+ return False
144
+
145
+ except requests.exceptions.RequestException as e:
146
+ # Handle network errors during polling gracefully
147
+ live_status.update("Waiting... (network error, retrying)")
148
+ console.print(f"\n[bold yellow]Warning:[/] Network error during polling: {e}. Retrying...")
149
+ # Optional: implement backoff strategy
150
+ time.sleep(interval * 2) # Simple backoff
151
+
152
+ except requests.exceptions.HTTPError as e: # Catch HTTPError specifically for handle_api_error
153
+ handle_api_error(e, console)
154
+ except requests.exceptions.RequestException as e: # Catch other request errors
155
+ console.print(f"\n[bold red]Network Error:[/] {e}")
156
+ return False
157
+ except Exception as e:
158
+ console.print(f"\n[bold red]An unexpected error occurred during login:[/] {e}")
159
+ return False
160
+
161
+
34
162
  def main() -> None:
35
163
  """Main function for the Weco CLI."""
164
+ # --- Argument Parsing ---
36
165
  parser = argparse.ArgumentParser(
37
166
  description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
38
167
  )
39
- parser.add_argument("--source", type=str, required=True, help="Path to the source code (e.g. optimize.py)")
40
- parser.add_argument(
168
+ # Add subparsers for commands like 'run' and 'logout'
169
+ subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True) # Make command required
170
+
171
+ # --- Run Command ---
172
+ run_parser = subparsers.add_parser(
173
+ "run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter
174
+ )
175
+ # Add arguments specific to the 'run' command to the run_parser
176
+ run_parser.add_argument("--source", type=str, required=True, help="Path to the source code file (e.g. optimize.py)")
177
+ run_parser.add_argument(
41
178
  "--eval-command", type=str, required=True, help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1')"
42
179
  )
43
- parser.add_argument("--metric", type=str, required=True, help="Metric to optimize")
44
- parser.add_argument(
180
+ run_parser.add_argument("--metric", type=str, required=True, help="Metric to optimize")
181
+ run_parser.add_argument(
45
182
  "--maximize",
46
183
  type=str,
47
184
  choices=["true", "false"],
48
185
  required=True,
49
186
  help="Specify 'true' to maximize the metric or 'false' to minimize.",
50
187
  )
51
- parser.add_argument("--steps", type=int, required=True, help="Number of steps to run")
52
- parser.add_argument("--model", type=str, required=True, help="Model to use for optimization")
53
- parser.add_argument("--log-dir", type=str, default=".runs", help="Directory to store logs and results")
54
- parser.add_argument(
188
+ run_parser.add_argument("--steps", type=int, required=True, help="Number of steps to run")
189
+ run_parser.add_argument("--model", type=str, required=True, help="Model to use for optimization")
190
+ run_parser.add_argument("--log-dir", type=str, default=".runs", help="Directory to store logs and results")
191
+ run_parser.add_argument(
55
192
  "--additional-instructions",
56
193
  default=None,
57
194
  type=str,
58
195
  help="Description of additional instruction or path to a file containing additional instructions",
59
196
  )
60
- parser.add_argument(
197
+ run_parser.add_argument(
61
198
  "--preserve-source",
62
199
  action="store_true",
63
200
  help="If set, do not overwrite the original source file; only save modified versions in the runs directory",
64
201
  )
202
+
203
+ # --- Logout Command ---
204
+ _ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
205
+
65
206
  args = parser.parse_args()
66
207
 
67
- try:
68
- with console.status("[bold green]Loading Modules..."):
69
- # Define optimization session config
208
+ # --- Handle Logout Command ---
209
+ if args.command == "logout":
210
+ clear_api_key()
211
+ sys.exit(0)
212
+
213
+ # --- Handle Run Command ---
214
+ elif args.command == "run":
215
+ # --- Check Authentication ---
216
+ weco_api_key = load_weco_api_key()
217
+ llm_api_keys = read_api_keys_from_env() # Read keys from client environment
218
+
219
+ if not weco_api_key:
220
+ login_choice = Prompt.ask(
221
+ "Log in to Weco to save run history or use anonymously? ([bold]L[/]ogin / [bold]S[/]kip)",
222
+ choices=["l", "s"],
223
+ default="s",
224
+ ).lower()
225
+
226
+ if login_choice == "l":
227
+ console.print("[cyan]Starting login process...[/]")
228
+ if not perform_login(console):
229
+ console.print("[bold red]Login process failed or was cancelled.[/]")
230
+ sys.exit(1)
231
+ weco_api_key = load_weco_api_key()
232
+ if not weco_api_key:
233
+ console.print("[bold red]Error: Login completed but failed to retrieve API key.[/]")
234
+ sys.exit(1)
235
+
236
+ elif login_choice == "s":
237
+ console.print("[yellow]Proceeding anonymously. LLM API keys must be provided via environment variables.[/]")
238
+ if not llm_api_keys:
239
+ console.print(
240
+ "[bold red]Error:[/] No LLM API keys found in environment (e.g., OPENAI_API_KEY). Cannot proceed anonymously."
241
+ )
242
+ sys.exit(1)
243
+
244
+ # --- Prepare API Call Arguments ---
245
+ auth_headers = {}
246
+
247
+ if weco_api_key:
248
+ auth_headers["Authorization"] = f"Bearer {weco_api_key}"
249
+ # Backend will decide whether to use client keys based on auth status
250
+
251
+ # --- Main Run Logic ---
252
+ try:
253
+ # --- Configuration Loading ---
70
254
  evaluation_command = args.eval_command
71
255
  metric_name = args.metric
72
256
  maximize = args.maximize == "true"
73
257
  steps = args.steps
74
258
  code_generator_config = {"model": args.model}
75
- evaluator_config = {"model": args.model}
259
+ evaluator_config = {
260
+ "model": args.model,
261
+ "include_analysis": False, # NOTE: False for now
262
+ }
76
263
  search_policy_config = {
77
- "num_drafts": max(1, math.ceil(0.15 * steps)), # 15% of steps
264
+ "num_drafts": max(1, math.ceil(0.15 * steps)),
78
265
  "debug_prob": 0.5,
79
- "max_debug_depth": max(1, math.ceil(0.1 * steps)), # 10% of steps
266
+ "max_debug_depth": max(1, math.ceil(0.1 * steps)),
80
267
  }
81
- # Read API keys
82
- api_keys = read_api_keys_from_env()
83
- # API request timeout
84
- timeout = 800
85
-
86
268
  # Read additional instructions
87
269
  additional_instructions = read_additional_instructions(additional_instructions=args.additional_instructions)
88
- # Read source code
270
+ # Read source code path
89
271
  source_fp = pathlib.Path(args.source)
272
+ # Read source code content
90
273
  source_code = read_from_path(fp=source_fp, is_json=False)
274
+ # API request timeout
275
+ timeout = 800
91
276
 
92
- # Initialize panels
93
- summary_panel = SummaryPanel(
94
- maximize=maximize, metric_name=metric_name, total_steps=steps, model=args.model, runs_dir=args.log_dir
95
- )
96
- plan_panel = PlanPanel()
97
- solution_panels = SolutionPanels(metric_name=metric_name, source_fp=source_fp)
98
- eval_output_panel = EvaluationOutputPanel()
99
- tree_panel = MetricTreePanel(maximize=maximize)
100
- layout = create_optimization_layout()
101
- end_optimization_layout = create_end_optimization_layout()
102
-
103
- # Start optimization session
104
- session_response = start_optimization_session(
105
- console=console,
106
- source_code=source_code,
107
- evaluation_command=evaluation_command,
108
- metric_name=metric_name,
109
- maximize=maximize,
110
- steps=steps,
111
- code_generator_config=code_generator_config,
112
- evaluator_config=evaluator_config,
113
- search_policy_config=search_policy_config,
114
- additional_instructions=additional_instructions,
115
- api_keys=api_keys,
116
- timeout=timeout,
117
- )
118
-
119
- # Define the refresh rate
120
- refresh_rate = 4
121
- with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
122
- # Define the runs directory (.runs/<session-id>)
123
- session_id = session_response["session_id"]
124
- runs_dir = pathlib.Path(args.log_dir) / session_id
125
- runs_dir.mkdir(parents=True, exist_ok=True)
126
-
127
- # Save the original code (.runs/<session-id>/original.<extension>)
128
- runs_copy_source_fp = runs_dir / f"original{source_fp.suffix}"
129
- write_to_path(fp=runs_copy_source_fp, content=source_code)
130
-
131
- # Write the code string to the source file path
132
- # Do this after the original code is saved
133
- if not args.preserve_source:
134
- write_to_path(fp=source_fp, content=session_response["code"])
135
-
136
- # Update the panels with the initial solution
137
- # Add session id now that we have it
138
- summary_panel.session_id = session_id
139
- # Set the step of the progress bar
140
- summary_panel.set_step(step=0)
141
- # Update the token counts
142
- summary_panel.update_token_counts(usage=session_response["usage"])
143
- # Update the plan
144
- plan_panel.update(plan=session_response["plan"])
145
- # Build the metric tree
146
- tree_panel.build_metric_tree(
147
- nodes=[
148
- {
149
- "solution_id": session_response["solution_id"],
150
- "parent_id": None,
151
- "code": session_response["code"],
152
- "step": 0,
153
- "metric_value": None,
154
- "is_buggy": False,
155
- }
156
- ]
157
- )
158
- # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
159
- tree_panel.set_unevaluated_node(node_id=session_response["solution_id"])
160
- # Update the solution panels with the initial solution and get the panel displays
161
- solution_panels.update(
162
- current_node=Node(
163
- id=session_response["solution_id"],
164
- parent_id=None,
165
- code=session_response["code"],
166
- metric=None,
167
- is_buggy=False,
168
- ),
169
- best_node=None,
277
+ # --- Panel Initialization ---
278
+ summary_panel = SummaryPanel(
279
+ maximize=maximize, metric_name=metric_name, total_steps=steps, model=args.model, runs_dir=args.log_dir
170
280
  )
171
- current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=0)
172
-
173
- # Update the entire layout
174
- smooth_update(
175
- live=live,
176
- layout=layout,
177
- sections_to_update=[
178
- ("summary", summary_panel.get_display()),
179
- ("plan", plan_panel.get_display()),
180
- ("tree", tree_panel.get_display()),
181
- ("current_solution", current_solution_panel),
182
- ("best_solution", best_solution_panel),
183
- ("eval_output", eval_output_panel.get_display()),
184
- ],
185
- transition_delay=0.1, # Slightly longer delay for initial display
281
+ plan_panel = PlanPanel()
282
+ solution_panels = SolutionPanels(metric_name=metric_name, source_fp=source_fp)
283
+ eval_output_panel = EvaluationOutputPanel()
284
+ tree_panel = MetricTreePanel(maximize=maximize)
285
+ layout = create_optimization_layout()
286
+ end_optimization_layout = create_end_optimization_layout()
287
+
288
+ # --- Start Optimization Session ---
289
+ session_response = start_optimization_session(
290
+ console=console,
291
+ source_code=source_code,
292
+ evaluation_command=evaluation_command,
293
+ metric_name=metric_name,
294
+ maximize=maximize,
295
+ steps=steps,
296
+ code_generator_config=code_generator_config,
297
+ evaluator_config=evaluator_config,
298
+ search_policy_config=search_policy_config,
299
+ additional_instructions=additional_instructions,
300
+ api_keys=llm_api_keys, # Pass client LLM keys
301
+ auth_headers=auth_headers, # Pass Weco key if logged in
302
+ timeout=timeout,
186
303
  )
187
304
 
188
- # Run evaluation on the initial solution
189
- term_out = run_evaluation(eval_command=args.eval_command)
305
+ # --- Live Update Loop ---
306
+ refresh_rate = 4
307
+ with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
308
+ # Define the runs directory (.runs/<session-id>)
309
+ session_id = session_response["session_id"]
310
+ runs_dir = pathlib.Path(args.log_dir) / session_id
311
+ runs_dir.mkdir(parents=True, exist_ok=True)
190
312
 
191
- # Update the evaluation output panel
192
- eval_output_panel.update(output=term_out)
193
- smooth_update(
194
- live=live,
195
- layout=layout,
196
- sections_to_update=[("eval_output", eval_output_panel.get_display())],
197
- transition_delay=0.1,
198
- )
313
+ # Save the original code (.runs/<session-id>/original.<extension>)
314
+ runs_copy_source_fp = runs_dir / f"original{source_fp.suffix}" # Use correct suffix
315
+ write_to_path(fp=runs_copy_source_fp, content=source_code)
199
316
 
200
- for step in range(1, steps):
201
- # Evaluate the current output and get the next solution
202
- eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
203
- console=console,
204
- session_id=session_id,
205
- execution_output=term_out,
206
- additional_instructions=None,
207
- api_keys=api_keys,
208
- timeout=timeout,
209
- )
210
- # Save next solution (.runs/<session-id>/step_<step>.<extension>)
211
- write_to_path(fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"])
212
-
213
- # Write the next solution to the source file
317
+ # Write the initial code string to the source file path (if not preserving)
214
318
  if not args.preserve_source:
215
- write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
319
+ write_to_path(fp=source_fp, content=session_response["code"])
216
320
 
217
- # Get the optimization session status for
218
- # the best solution, its score, and the history to plot the tree
219
- status_response = get_optimization_session_status(
220
- console=console, session_id=session_id, include_history=True, timeout=timeout
221
- )
222
-
223
- # Update the step of the progress bar
224
- summary_panel.set_step(step=step)
321
+ # Update the panels with the initial solution
322
+ summary_panel.set_session_id(session_id=session_id) # Add session id now that we have it
323
+ # Set the step of the progress bar
324
+ summary_panel.set_step(step=0)
225
325
  # Update the token counts
226
- summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
326
+ summary_panel.update_token_counts(usage=session_response["usage"])
227
327
  # Update the plan
228
- plan_panel.update(plan=eval_and_next_solution_response["plan"])
328
+ plan_panel.update(plan=session_response["plan"])
229
329
  # Build the metric tree
230
- tree_panel.build_metric_tree(nodes=status_response["history"])
330
+ tree_panel.build_metric_tree(
331
+ nodes=[
332
+ {
333
+ "solution_id": session_response["solution_id"],
334
+ "parent_id": None,
335
+ "code": session_response["code"],
336
+ "step": 0,
337
+ "metric_value": None,
338
+ "is_buggy": False,
339
+ }
340
+ ]
341
+ )
231
342
  # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
232
- tree_panel.set_unevaluated_node(node_id=eval_and_next_solution_response["solution_id"])
233
-
234
- # Update the solution panels with the next solution and best solution (and score)
235
- # Figure out if we have a best solution so far
236
- if status_response["best_result"] is not None:
237
- best_solution_node = Node(
238
- id=status_response["best_result"]["solution_id"],
239
- parent_id=status_response["best_result"]["parent_id"],
240
- code=status_response["best_result"]["code"],
241
- metric=status_response["best_result"]["metric_value"],
242
- is_buggy=status_response["best_result"]["is_buggy"],
243
- )
244
- else:
245
- best_solution_node = None
246
-
247
- # Create a node for the current solution
248
- current_solution_node = None
249
- for node in status_response["history"]:
250
- if node["solution_id"] == eval_and_next_solution_response["solution_id"]:
251
- current_solution_node = Node(
252
- id=node["solution_id"],
253
- parent_id=node["parent_id"],
254
- code=node["code"],
255
- metric=node["metric_value"],
256
- is_buggy=node["is_buggy"],
257
- )
258
- if current_solution_node is None:
259
- raise ValueError("Current solution node not found in history")
260
- # Update the solution panels with the current and best solution
261
- solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
262
- current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
263
-
264
- # Clear evaluation output since we are running a evaluation on a new solution
265
- eval_output_panel.clear()
266
-
267
- # Update displays with smooth transitions
343
+ tree_panel.set_unevaluated_node(node_id=session_response["solution_id"])
344
+ # Update the solution panels with the initial solution and get the panel displays
345
+ solution_panels.update(
346
+ current_node=Node(
347
+ id=session_response["solution_id"],
348
+ parent_id=None,
349
+ code=session_response["code"],
350
+ metric=None,
351
+ is_buggy=False,
352
+ ),
353
+ best_node=None,
354
+ )
355
+ current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=0)
356
+ # Update the live layout with the initial solution panels
268
357
  smooth_update(
269
358
  live=live,
270
359
  layout=layout,
271
360
  sections_to_update=[
272
361
  ("summary", summary_panel.get_display()),
273
362
  ("plan", plan_panel.get_display()),
274
- ("tree", tree_panel.get_display()),
363
+ ("tree", tree_panel.get_display(is_done=False)),
275
364
  ("current_solution", current_solution_panel),
276
365
  ("best_solution", best_solution_panel),
277
366
  ("eval_output", eval_output_panel.get_display()),
278
367
  ],
279
- transition_delay=0.08, # Slightly longer delay for more noticeable transitions
368
+ transition_delay=0.1,
280
369
  )
281
370
 
282
- # Run evaluation on the current solution
371
+ # Run evaluation on the initial solution
283
372
  term_out = run_evaluation(eval_command=args.eval_command)
284
- eval_output_panel.update(output=term_out)
285
373
 
286
- # Update evaluation output with a smooth transition
374
+ # Update the evaluation output panel
375
+ eval_output_panel.update(output=term_out)
287
376
  smooth_update(
288
377
  live=live,
289
378
  layout=layout,
290
379
  sections_to_update=[("eval_output", eval_output_panel.get_display())],
291
- transition_delay=0.1, # Slightly longer delay for evaluation results
380
+ transition_delay=0.1,
292
381
  )
293
382
 
294
- # Ensure we pass evaluation results for the last step's generated solution
295
- eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
296
- console=console,
297
- session_id=session_id,
298
- execution_output=term_out,
299
- additional_instructions=None,
300
- api_keys=api_keys,
301
- timeout=timeout,
302
- )
383
+ for step in range(1, steps):
384
+ # Re-read instructions from the original source (file path or string) BEFORE each suggest call
385
+ current_additional_instructions = read_additional_instructions(
386
+ additional_instructions=args.additional_instructions
387
+ )
303
388
 
304
- # Update the progress bar
305
- summary_panel.set_step(step=steps)
306
- # Update the token counts
307
- summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
308
- # No need to update the plan panel since we have finished the optimization
309
- # Get the optimization session status for
310
- # the best solution, its score, and the history to plot the tree
311
- status_response = get_optimization_session_status(
312
- console=console, session_id=session_id, include_history=True, timeout=timeout
313
- )
314
- # Build the metric tree
315
- tree_panel.build_metric_tree(nodes=status_response["history"])
316
- # No need to set any solution to unevaluated since we have finished the optimization
317
- # and all solutions have been evaluated
318
- # No neeed to update the current solution panel since we have finished the optimization
319
- # We only need to update the best solution panel
320
- # Figure out if we have a best solution so far
321
- if status_response["best_result"] is not None:
322
- best_solution_node = Node(
323
- id=status_response["best_result"]["solution_id"],
324
- parent_id=status_response["best_result"]["parent_id"],
325
- code=status_response["best_result"]["code"],
326
- metric=status_response["best_result"]["metric_value"],
327
- is_buggy=status_response["best_result"]["is_buggy"],
389
+ # Send feedback and get next suggestion
390
+ eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
391
+ console=console,
392
+ session_id=session_id,
393
+ execution_output=term_out,
394
+ additional_instructions=current_additional_instructions, # Pass current instructions
395
+ api_keys=llm_api_keys, # Pass client LLM keys
396
+ auth_headers=auth_headers, # Pass Weco key if logged in
397
+ timeout=timeout,
398
+ )
399
+ # Save next solution (.runs/<session-id>/step_<step>.<extension>)
400
+ write_to_path(
401
+ fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"]
402
+ )
403
+
404
+ # Write the next solution to the source file
405
+ if not args.preserve_source:
406
+ write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
407
+
408
+ # Get the optimization session status for
409
+ # the best solution, its score, and the history to plot the tree
410
+ status_response = get_optimization_session_status(
411
+ console=console,
412
+ session_id=session_id,
413
+ include_history=True,
414
+ timeout=timeout,
415
+ auth_headers=auth_headers,
416
+ )
417
+
418
+ # Update the step of the progress bar
419
+ summary_panel.set_step(step=step)
420
+ # Update the token counts
421
+ summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
422
+ # Update the plan
423
+ plan_panel.update(plan=eval_and_next_solution_response["plan"])
424
+ # Build the metric tree
425
+ tree_panel.build_metric_tree(nodes=status_response["history"])
426
+ # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
427
+ tree_panel.set_unevaluated_node(node_id=eval_and_next_solution_response["solution_id"])
428
+
429
+ # Update the solution panels with the next solution and best solution (and score)
430
+ # Figure out if we have a best solution so far
431
+ if status_response["best_result"] is not None:
432
+ best_solution_node = Node(
433
+ id=status_response["best_result"]["solution_id"],
434
+ parent_id=status_response["best_result"]["parent_id"],
435
+ code=status_response["best_result"]["code"],
436
+ metric=status_response["best_result"]["metric_value"],
437
+ is_buggy=status_response["best_result"]["is_buggy"],
438
+ )
439
+ else:
440
+ best_solution_node = None
441
+
442
+ # Create a node for the current solution
443
+ current_solution_node = None
444
+ for node in status_response["history"]:
445
+ if node["solution_id"] == eval_and_next_solution_response["solution_id"]:
446
+ current_solution_node = Node(
447
+ id=node["solution_id"],
448
+ parent_id=node["parent_id"],
449
+ code=node["code"],
450
+ metric=node["metric_value"],
451
+ is_buggy=node["is_buggy"],
452
+ )
453
+ if current_solution_node is None:
454
+ raise ValueError("Current solution node not found in history")
455
+ # Update the solution panels with the current and best solution
456
+ solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
457
+ current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
458
+
459
+ # Clear evaluation output since we are running a evaluation on a new solution
460
+ eval_output_panel.clear()
461
+
462
+ # Update displays with smooth transitions
463
+ smooth_update(
464
+ live=live,
465
+ layout=layout,
466
+ sections_to_update=[
467
+ ("summary", summary_panel.get_display()),
468
+ ("plan", plan_panel.get_display()),
469
+ ("tree", tree_panel.get_display(is_done=False)),
470
+ ("current_solution", current_solution_panel),
471
+ ("best_solution", best_solution_panel),
472
+ ("eval_output", eval_output_panel.get_display()),
473
+ ],
474
+ transition_delay=0.08, # Slightly longer delay for more noticeable transitions
475
+ )
476
+
477
+ # Run evaluation on the current solution
478
+ term_out = run_evaluation(eval_command=args.eval_command)
479
+ eval_output_panel.update(output=term_out)
480
+
481
+ # Update evaluation output with a smooth transition
482
+ smooth_update(
483
+ live=live,
484
+ layout=layout,
485
+ sections_to_update=[("eval_output", eval_output_panel.get_display())],
486
+ transition_delay=0.1, # Slightly longer delay for evaluation results
487
+ )
488
+
489
+ # Re-read instructions from the original source (file path or string) BEFORE each suggest call
490
+ current_additional_instructions = read_additional_instructions(
491
+ additional_instructions=args.additional_instructions
328
492
  )
329
- else:
330
- best_solution_node = None
331
- solution_panels.update(current_node=None, best_node=best_solution_node)
332
- _, best_solution_panel = solution_panels.get_display(current_step=steps)
333
-
334
- # Update the end optimization layout
335
- final_message = (
336
- f"{summary_panel.metric_name.capitalize()} {'maximized' if summary_panel.maximize else 'minimized'}! Best solution {summary_panel.metric_name.lower()} = [green]{status_response['best_result']['metric_value']}[/] 🏆"
337
- if best_solution_node is not None and best_solution_node.metric is not None
338
- else "[red] No valid solution found.[/]"
339
- )
340
- end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
341
- end_optimization_layout["tree"].update(tree_panel.get_display())
342
- end_optimization_layout["best_solution"].update(best_solution_panel)
343
-
344
- # Save optimization results
345
- # If the best solution does not exist or is has not been measured at the end of the optimization
346
- # save the original solution as the best solution
347
- if best_solution_node is not None:
348
- best_solution_code = best_solution_node.code
349
- best_solution_score = best_solution_node.metric
350
- else:
351
- best_solution_code = None
352
- best_solution_score = None
353
-
354
- if best_solution_code is None or best_solution_score is None:
355
- best_solution_content = (
356
- f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_copy_source_fp, is_json=False)}"
493
+
494
+ # Ensure we pass evaluation results for the last step's generated solution
495
+ eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
496
+ console=console,
497
+ session_id=session_id,
498
+ execution_output=term_out,
499
+ additional_instructions=current_additional_instructions,
500
+ api_keys=llm_api_keys,
501
+ timeout=timeout,
502
+ auth_headers=auth_headers,
357
503
  )
358
- else:
359
- # Format score for the comment
360
- best_score_str = (
361
- format_number(best_solution_score)
362
- if best_solution_score is not None and isinstance(best_solution_score, (int, float))
363
- else "N/A"
504
+
505
+ # Update the progress bar
506
+ summary_panel.set_step(step=steps)
507
+ # Update the token counts
508
+ summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
509
+ # No need to update the plan panel since we have finished the optimization
510
+ # Get the optimization session status for
511
+ # the best solution, its score, and the history to plot the tree
512
+ status_response = get_optimization_session_status(
513
+ console=console, session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
514
+ )
515
+ # Build the metric tree
516
+ tree_panel.build_metric_tree(nodes=status_response["history"])
517
+ # No need to set any solution to unevaluated since we have finished the optimization
518
+ # and all solutions have been evaluated
519
+ # No neeed to update the current solution panel since we have finished the optimization
520
+ # We only need to update the best solution panel
521
+ # Figure out if we have a best solution so far
522
+ if status_response["best_result"] is not None:
523
+ best_solution_node = Node(
524
+ id=status_response["best_result"]["solution_id"],
525
+ parent_id=status_response["best_result"]["parent_id"],
526
+ code=status_response["best_result"]["code"],
527
+ metric=status_response["best_result"]["metric_value"],
528
+ is_buggy=status_response["best_result"]["is_buggy"],
529
+ )
530
+ else:
531
+ best_solution_node = None
532
+ solution_panels.update(current_node=None, best_node=best_solution_node)
533
+ _, best_solution_panel = solution_panels.get_display(current_step=steps)
534
+
535
+ # Update the end optimization layout
536
+ final_message = (
537
+ f"{summary_panel.metric_name.capitalize()} {'maximized' if summary_panel.maximize else 'minimized'}! Best solution {summary_panel.metric_name.lower()} = [green]{status_response['best_result']['metric_value']}[/] 🏆"
538
+ if best_solution_node is not None and best_solution_node.metric is not None
539
+ else "[red] No valid solution found.[/]"
364
540
  )
365
- best_solution_content = f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
541
+ end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
542
+ end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
543
+ end_optimization_layout["best_solution"].update(best_solution_panel)
544
+
545
+ # Save optimization results
546
+ # If the best solution does not exist or is has not been measured at the end of the optimization
547
+ # save the original solution as the best solution
548
+ if best_solution_node is not None:
549
+ best_solution_code = best_solution_node.code
550
+ best_solution_score = best_solution_node.metric
551
+ else:
552
+ best_solution_code = None
553
+ best_solution_score = None
366
554
 
367
- # Save best solution to .runs/<session-id>/best.<extension>
368
- write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
555
+ if best_solution_code is None or best_solution_score is None:
556
+ best_solution_content = (
557
+ f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_copy_source_fp, is_json=False)}"
558
+ )
559
+ else:
560
+ # Format score for the comment
561
+ best_score_str = (
562
+ format_number(best_solution_score)
563
+ if best_solution_score is not None and isinstance(best_solution_score, (int, float))
564
+ else "N/A"
565
+ )
566
+ best_solution_content = (
567
+ f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
568
+ )
369
569
 
370
- # write the best solution to the source file
371
- if not args.preserve_source:
372
- write_to_path(fp=source_fp, content=best_solution_content)
570
+ # Save best solution to .runs/<session-id>/best.<extension>
571
+ write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
373
572
 
374
- console.print(end_optimization_layout)
573
+ # write the best solution to the source file
574
+ if not args.preserve_source:
575
+ write_to_path(fp=source_fp, content=best_solution_content)
375
576
 
376
- except Exception as e:
377
- console.print(Panel(f"[bold red]Error: {str(e)}", title="[bold red]Error", border_style="red"))
378
- sys.exit(1)
577
+ console.print(end_optimization_layout)
578
+
579
+ except Exception as e:
580
+ console.print(Panel(f"[bold red]Error: {str(e)}", title="[bold red]Error", border_style="red"))
581
+ # Print traceback for debugging
582
+ console.print_exception(show_locals=True)
583
+ sys.exit(1)
weco/panels.py CHANGED
@@ -7,6 +7,7 @@ from rich.syntax import Syntax
7
7
  from typing import Dict, List, Optional, Union, Tuple
8
8
  from .utils import format_number
9
9
  import pathlib
10
+ from .__init__ import __dashboard_url__
10
11
 
11
12
 
12
13
  class SummaryPanel:
@@ -22,6 +23,7 @@ class SummaryPanel:
22
23
  self.model = model
23
24
  self.runs_dir = runs_dir
24
25
  self.session_id = session_id if session_id is not None else "N/A"
26
+ self.dashboard_url = "N/A"
25
27
  self.progress = Progress(
26
28
  TextColumn("[progress.description]{task.description}"),
27
29
  BarColumn(bar_width=20),
@@ -32,6 +34,15 @@ class SummaryPanel:
32
34
  )
33
35
  self.task_id = self.progress.add_task("", total=total_steps)
34
36
 
37
+ def set_session_id(self, session_id: str):
38
+ """Set the session ID."""
39
+ self.session_id = session_id
40
+ self.set_dashboard_url(session_id=session_id)
41
+
42
+ def set_dashboard_url(self, session_id: str):
43
+ """Set the dashboard URL."""
44
+ self.dashboard_url = f"{__dashboard_url__}/runs/{session_id}"
45
+
35
46
  def set_step(self, step: int):
36
47
  """Set the current step."""
37
48
  self.progress.update(self.task_id, completed=step)
@@ -61,6 +72,9 @@ class SummaryPanel:
61
72
  # Log directory
62
73
  summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.session_id}[/]")
63
74
  summary_table.add_row("")
75
+ # Dashboard link
76
+ summary_table.add_row(f"[bold cyan]Dashboard:[/] [blue underline]{self.dashboard_url}[/]")
77
+ summary_table.add_row("")
64
78
  # Token counts
65
79
  summary_table.add_row(
66
80
  f"[bold cyan]Tokens:[/] ↑[yellow]{format_number(self.total_input_tokens)}[/] ↓[yellow]{format_number(self.total_output_tokens)}[/] = [green]{format_number(self.total_input_tokens + self.total_output_tokens)}[/]"
@@ -229,11 +243,15 @@ class MetricTreePanel:
229
243
 
230
244
  return tree
231
245
 
232
- def get_display(self) -> Panel:
246
+ def get_display(self, is_done: bool) -> Panel:
233
247
  """Get a panel displaying the solution tree."""
234
248
  # Make sure the metric tree is built before calling build_rich_tree
235
249
  return Panel(
236
- self._build_rich_tree(), title="[bold]🔎 Exploring Solutions...", border_style="green", expand=True, padding=(0, 1)
250
+ self._build_rich_tree(),
251
+ title="[bold]🔎 Exploring Solutions..." if not is_done else "[bold]🔎 Optimization Complete!",
252
+ border_style="green",
253
+ expand=True,
254
+ padding=(0, 1),
237
255
  )
238
256
 
239
257
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: weco
3
- Version: 0.2.12
3
+ Version: 0.2.14
4
4
  Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
5
5
  Author-email: Weco AI Team <contact@weco.ai>
6
6
  License: MIT
@@ -32,7 +32,7 @@ Example applications include:
32
32
 
33
33
  - **GPU Kernel Optimization**: Reimplement PyTorch functions using CUDA, Triton or Metal, optimizing for `latency`, `throughput`, or `memory_bandwidth`.
34
34
  - **Model Development**: Tune feature transformations or architectures, optimizing for `validation_accuracy`, `AUC`, or `Sharpe Ratio`.
35
- - **Prompt Engineering**: Refine prompts for LLMs, optimizing for `win_rate`, `relevance`, or `format_adherence`
35
+ - **Prompt Engineering**: Refine prompts for LLMs, optimizing for `win_rate`, `relevance`, or `format_adherence`
36
36
 
37
37
  https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
38
38
 
@@ -42,7 +42,7 @@ https://github.com/user-attachments/assets/cb724ef1-bff6-4757-b457-d3b2201ede81
42
42
 
43
43
  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.
44
44
 
45
- ![image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
45
+ [image](https://github.com/user-attachments/assets/a6ed63fa-9c40-498e-aa98-a873e5786509)
46
46
 
47
47
  ---
48
48
 
@@ -54,17 +54,38 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
54
54
  pip install weco
55
55
  ```
56
56
 
57
- 2. **Configure API Keys:**
57
+ 2. **Set Up LLM API Keys (Required):**
58
58
 
59
- Set the appropriate environment variables for your desired language model provider:
59
+ `weco` requires API keys for the Large Language Models (LLMs) it uses internally. You **must** provide these keys via environment variables:
60
60
 
61
- - **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
62
- - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
63
- - **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.)
61
+ - **OpenAI:** `export OPENAI_API_KEY="your_key_here"`
62
+ - **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
63
+ - **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.)
64
+
65
+ The optimization process will fail if the necessary keys for the chosen model are not found in your environment.
66
+
67
+ 3. **Log In to Weco (Optional):**
68
+
69
+ 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:
70
+
71
+ - When you first run `weco run`, you'll be prompted if you want to log in or proceed anonymously.
72
+ - 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.
73
+ - You then authenticate in the browser. Once authenticated, the CLI will detect this and complete the login.
74
+ - This saves a Weco-specific API key locally (typically at `~/.config/weco/credentials.json`).
75
+
76
+ 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.
77
+
78
+ To log out and remove your saved Weco API key, use the `weco logout` command.
64
79
 
65
80
  ---
66
81
 
67
82
  ## Usage
83
+
84
+ The CLI has two main commands:
85
+
86
+ - `weco run`: Initiates the code optimization process.
87
+ - `weco logout`: Logs you out of your Weco account.
88
+
68
89
  <div style="background-color: #fff3cd; border: 1px solid #ffeeba; padding: 15px; border-radius: 4px; margin-bottom: 15px;">
69
90
  <strong>⚠️ Warning: Code Modification</strong><br>
70
91
  <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.
@@ -72,7 +93,11 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
72
93
 
73
94
  ---
74
95
 
75
- ### Example: Optimizing Simple PyTorch Operations
96
+ ### `weco run` Command
97
+
98
+ This command starts the optimization process.
99
+
100
+ **Example: Optimizing Simple PyTorch Operations**
76
101
 
77
102
  This basic example shows how to optimize a simple PyTorch function for speedup.
78
103
 
@@ -86,7 +111,7 @@ cd examples/hello-kernel-world
86
111
  pip install torch
87
112
 
88
113
  # Run Weco
89
- weco --source optimize.py \
114
+ weco run --source optimize.py \
90
115
  --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
91
116
  --metric speedup \
92
117
  --maximize true \
@@ -99,19 +124,29 @@ weco --source optimize.py \
99
124
 
100
125
  ---
101
126
 
102
- ### Command Line Arguments
127
+ **Arguments for `weco run`:**
103
128
 
104
- | Argument | Description | Required |
105
- | :-------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :------- |
106
- | `--source` | Path to the source code file that will be optimized (e.g., `optimize.py`). | Yes |
107
- | `--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 |
108
- | `--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 |
109
- | `--maximize` | Whether to maximize (`true`) or minimize (`false`) the metric. | Yes |
110
- | `--steps` | Number of optimization steps (LLM iterations) to run. | Yes |
111
- | `--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 |
112
- | `--additional-instructions` | (Optional) Natural language description of specific instructions OR path to a file containing detailed instructions to guide the LLM. | No |
113
- | `--log-dir` | (Optional) Path to the directory to log intermediate steps and final optimization result. Defaults to `.runs/`. | No |
114
- | `--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 |
129
+ | Argument | Description | Required |
130
+ | :-------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | :------- |
131
+ | `--source` | Path to the source code file that will be optimized (e.g., `optimize.py`). | Yes |
132
+ | `--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 |
133
+ | `--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 |
134
+ | `--maximize` | Whether to maximize (`true`) or minimize (`false`) the metric. | Yes |
135
+ | `--steps` | Number of optimization steps (LLM iterations) to run. | Yes |
136
+ | `--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 |
137
+ | `--additional-instructions` | (Optional) Natural language description of specific instructions OR path to a file containing detailed instructions to guide the LLM. | No |
138
+ | `--log-dir` | (Optional) Path to the directory to log intermediate steps and final optimization result. Defaults to `.runs/`. | No |
139
+ | `--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 |
140
+
141
+ ---
142
+
143
+ ### `weco logout` Command
144
+
145
+ This command logs you out by removing the locally stored Weco API key.
146
+
147
+ ```bash
148
+ weco logout
149
+ ```
115
150
 
116
151
  ---
117
152
 
@@ -120,6 +155,7 @@ weco --source optimize.py \
120
155
  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.
121
156
 
122
157
  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.
158
+
123
159
  <p align="center">
124
160
  <img src="https://github.com/user-attachments/assets/ff0e471d-2f50-4e2d-b718-874862f533df" alt="RE-Bench Performance Across Time" width="60%"/>
125
161
  </p>
@@ -146,23 +182,25 @@ Final speedup value = 1.5
146
182
 
147
183
  Weco will parse this output to extract the numerical value (1.5 in this case) associated with the metric name ('speedup').
148
184
 
149
-
150
185
  ## Contributing
151
186
 
152
187
  We welcome contributions! To get started:
153
188
 
154
189
  1. **Fork and Clone the Repository:**
190
+
155
191
  ```bash
156
192
  git clone https://github.com/WecoAI/weco-cli.git
157
193
  cd weco-cli
158
194
  ```
159
195
 
160
196
  2. **Install Development Dependencies:**
197
+
161
198
  ```bash
162
199
  pip install -e ".[dev]"
163
200
  ```
164
201
 
165
202
  3. **Create a Feature Branch:**
203
+
166
204
  ```bash
167
205
  git checkout -b feature/your-feature-name
168
206
  ```
@@ -0,0 +1,12 @@
1
+ weco/__init__.py,sha256=q4zeQ8CJq8NJyRcBMmST6zCzmK_HjNSFChrNBzJ9oks,426
2
+ weco/api.py,sha256=z2DCe0kQJaaBHo-Vml52GRc9nwQBnrKATXoY3UfPojw,3842
3
+ weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
4
+ weco/cli.py,sha256=7pGJiyoBO7n6Xocwusx9iVqsaPj8OuPaJ1A2ks83Ekw,28961
5
+ weco/panels.py,sha256=gB4rZbCvqzewUCBcILvyyU4fnOQLwFgHCGmtn-ZlgSo,13385
6
+ weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
7
+ weco-0.2.14.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
8
+ weco-0.2.14.dist-info/METADATA,sha256=855avaVBC-wUyFDWIvITw0t4o_v4s90Us-t_3lQnPGw,10851
9
+ weco-0.2.14.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
10
+ weco-0.2.14.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
11
+ weco-0.2.14.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
12
+ weco-0.2.14.dist-info/RECORD,,
@@ -1,11 +0,0 @@
1
- weco/__init__.py,sha256=6mgJEq6p8NTpOgnccei6V4oZiwYU57om8m1g4-GW4BQ,125
2
- weco/api.py,sha256=89lB2572jApAxkA0DDppDnJKBwvZTa3kH9jFpC0LFDQ,3313
3
- weco/cli.py,sha256=TP1pAFjJNHO3rVJvpXegwWFu7oBUZnVSyJfaHgqTPYs,17920
4
- weco/panels.py,sha256=R_df-VAbWyLoqCA9A6UzbIGZ9sm2IgJO4idnyjmrHQk,12701
5
- weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
6
- weco-0.2.12.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
7
- weco-0.2.12.dist-info/METADATA,sha256=nSMO2Wa7ijmH_8u6C-TibsxpMFxooTso7o8j0TEL2x4,9379
8
- weco-0.2.12.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
9
- weco-0.2.12.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
10
- weco-0.2.12.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
11
- weco-0.2.12.dist-info/RECORD,,
File without changes