weco 0.2.16__py3-none-any.whl → 0.2.18__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 +2 -1
- weco/api.py +117 -41
- weco/cli.py +165 -23
- weco/panels.py +1 -1
- weco/utils.py +32 -0
- {weco-0.2.16.dist-info → weco-0.2.18.dist-info}/METADATA +19 -40
- weco-0.2.18.dist-info/RECORD +12 -0
- {weco-0.2.16.dist-info → weco-0.2.18.dist-info}/WHEEL +1 -1
- weco-0.2.16.dist-info/RECORD +0 -12
- {weco-0.2.16.dist-info → weco-0.2.18.dist-info}/entry_points.txt +0 -0
- {weco-0.2.16.dist-info → weco-0.2.18.dist-info}/licenses/LICENSE +0 -0
- {weco-0.2.16.dist-info → weco-0.2.18.dist-info}/top_level.txt +0 -0
weco/__init__.py
CHANGED
weco/api.py
CHANGED
|
@@ -1,14 +1,20 @@
|
|
|
1
|
-
from typing import Dict, Any
|
|
1
|
+
from typing import Dict, Any, Optional
|
|
2
2
|
import rich
|
|
3
3
|
import requests
|
|
4
4
|
from weco import __pkg_version__, __base_url__
|
|
5
5
|
import sys
|
|
6
|
+
from rich.console import Console
|
|
6
7
|
|
|
7
8
|
|
|
8
9
|
def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
|
|
9
10
|
"""Extract and display error messages from API responses in a structured format."""
|
|
10
|
-
|
|
11
|
-
|
|
11
|
+
try:
|
|
12
|
+
detail = e.response.json()["detail"]
|
|
13
|
+
except (ValueError, KeyError): # Handle cases where response is not JSON or detail key is missing
|
|
14
|
+
detail = f"HTTP {e.response.status_code} Error: {e.response.text}"
|
|
15
|
+
console.print(f"[bold red]{detail}[/]")
|
|
16
|
+
# Avoid exiting here, let the caller decide if the error is fatal
|
|
17
|
+
# sys.exit(1)
|
|
12
18
|
|
|
13
19
|
|
|
14
20
|
def start_optimization_session(
|
|
@@ -28,25 +34,32 @@ def start_optimization_session(
|
|
|
28
34
|
) -> Dict[str, Any]:
|
|
29
35
|
"""Start the optimization session."""
|
|
30
36
|
with console.status("[bold green]Starting Optimization..."):
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
"
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
37
|
+
try:
|
|
38
|
+
response = requests.post(
|
|
39
|
+
f"{__base_url__}/sessions", # Path is relative to base_url
|
|
40
|
+
json={
|
|
41
|
+
"source_code": source_code,
|
|
42
|
+
"additional_instructions": additional_instructions,
|
|
43
|
+
"objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
|
|
44
|
+
"optimizer": {
|
|
45
|
+
"steps": steps,
|
|
46
|
+
"code_generator": code_generator_config,
|
|
47
|
+
"evaluator": evaluator_config,
|
|
48
|
+
"search_policy": search_policy_config,
|
|
49
|
+
},
|
|
50
|
+
"metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
|
|
42
51
|
},
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
52
|
+
headers=auth_headers, # Add headers
|
|
53
|
+
timeout=timeout,
|
|
54
|
+
)
|
|
55
|
+
response.raise_for_status()
|
|
56
|
+
return response.json()
|
|
57
|
+
except requests.exceptions.HTTPError as e:
|
|
58
|
+
handle_api_error(e, console)
|
|
59
|
+
sys.exit(1) # Exit if starting session fails
|
|
60
|
+
except requests.exceptions.RequestException as e:
|
|
61
|
+
console.print(f"[bold red]Network Error starting session: {e}[/]")
|
|
62
|
+
sys.exit(1)
|
|
50
63
|
|
|
51
64
|
|
|
52
65
|
def evaluate_feedback_then_suggest_next_solution(
|
|
@@ -58,29 +71,92 @@ def evaluate_feedback_then_suggest_next_solution(
|
|
|
58
71
|
timeout: int = 800,
|
|
59
72
|
) -> Dict[str, Any]:
|
|
60
73
|
"""Evaluate the feedback and suggest the next solution."""
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
74
|
+
try:
|
|
75
|
+
response = requests.post(
|
|
76
|
+
f"{__base_url__}/sessions/{session_id}/suggest", # Path is relative to base_url
|
|
77
|
+
json={
|
|
78
|
+
"execution_output": execution_output,
|
|
79
|
+
"additional_instructions": additional_instructions,
|
|
80
|
+
"metadata": {**api_keys},
|
|
81
|
+
},
|
|
82
|
+
headers=auth_headers, # Add headers
|
|
83
|
+
timeout=timeout,
|
|
84
|
+
)
|
|
85
|
+
response.raise_for_status()
|
|
86
|
+
return response.json()
|
|
87
|
+
except requests.exceptions.HTTPError as e:
|
|
88
|
+
# Allow caller to handle suggest errors, maybe retry or terminate
|
|
89
|
+
handle_api_error(e, Console()) # Use default console if none passed
|
|
90
|
+
raise # Re-raise the exception
|
|
91
|
+
except requests.exceptions.RequestException as e:
|
|
92
|
+
print(f"Network Error during suggest: {e}") # Use print as console might not be available
|
|
93
|
+
raise # Re-raise the exception
|
|
73
94
|
|
|
74
95
|
|
|
75
96
|
def get_optimization_session_status(
|
|
76
97
|
session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
|
|
77
98
|
) -> Dict[str, Any]:
|
|
78
99
|
"""Get the current status of the optimization session."""
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
100
|
+
try:
|
|
101
|
+
response = requests.get(
|
|
102
|
+
f"{__base_url__}/sessions/{session_id}", # Path is relative to base_url
|
|
103
|
+
params={"include_history": include_history},
|
|
104
|
+
headers=auth_headers,
|
|
105
|
+
timeout=timeout,
|
|
106
|
+
)
|
|
107
|
+
response.raise_for_status()
|
|
108
|
+
return response.json()
|
|
109
|
+
except requests.exceptions.HTTPError as e:
|
|
110
|
+
handle_api_error(e, Console()) # Use default console
|
|
111
|
+
raise # Re-raise
|
|
112
|
+
except requests.exceptions.RequestException as e:
|
|
113
|
+
print(f"Network Error getting status: {e}")
|
|
114
|
+
raise # Re-raise
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def send_heartbeat(
|
|
118
|
+
session_id: str,
|
|
119
|
+
auth_headers: dict = {},
|
|
120
|
+
timeout: int = 10, # Shorter timeout for non-critical heartbeat
|
|
121
|
+
) -> bool:
|
|
122
|
+
"""Send a heartbeat signal to the backend."""
|
|
123
|
+
try:
|
|
124
|
+
response = requests.put(f"{__base_url__}/sessions/{session_id}/heartbeat", headers=auth_headers, timeout=timeout)
|
|
125
|
+
response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
|
|
126
|
+
return True
|
|
127
|
+
except requests.exceptions.HTTPError as e:
|
|
128
|
+
# Log non-critical errors like 409 Conflict (session not running)
|
|
129
|
+
if e.response.status_code == 409:
|
|
130
|
+
print(f"Heartbeat ignored: Session {session_id} is not running.", file=sys.stderr)
|
|
131
|
+
else:
|
|
132
|
+
print(f"Heartbeat failed for session {session_id}: HTTP {e.response.status_code}", file=sys.stderr)
|
|
133
|
+
# Don't exit, just report failure
|
|
134
|
+
return False
|
|
135
|
+
except requests.exceptions.RequestException as e:
|
|
136
|
+
# Network errors are also non-fatal for heartbeats
|
|
137
|
+
print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
|
|
138
|
+
return False
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def report_termination(
|
|
142
|
+
session_id: str,
|
|
143
|
+
status_update: str,
|
|
144
|
+
reason: str,
|
|
145
|
+
details: Optional[str] = None,
|
|
146
|
+
auth_headers: dict = {},
|
|
147
|
+
timeout: int = 30, # Reasonably longer timeout for important termination message
|
|
148
|
+
) -> bool:
|
|
149
|
+
"""Report the termination reason to the backend."""
|
|
150
|
+
try:
|
|
151
|
+
response = requests.post(
|
|
152
|
+
f"{__base_url__}/sessions/{session_id}/terminate",
|
|
153
|
+
json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
|
|
154
|
+
headers=auth_headers,
|
|
155
|
+
timeout=timeout,
|
|
156
|
+
)
|
|
157
|
+
response.raise_for_status()
|
|
158
|
+
return True
|
|
159
|
+
except requests.exceptions.RequestException as e:
|
|
160
|
+
# Log failure, but don't prevent CLI exit
|
|
161
|
+
print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
|
|
162
|
+
return False
|
weco/cli.py
CHANGED
|
@@ -5,6 +5,9 @@ import math
|
|
|
5
5
|
import time
|
|
6
6
|
import requests
|
|
7
7
|
import webbrowser
|
|
8
|
+
import threading
|
|
9
|
+
import signal
|
|
10
|
+
import traceback
|
|
8
11
|
from rich.console import Console
|
|
9
12
|
from rich.live import Live
|
|
10
13
|
from rich.panel import Panel
|
|
@@ -15,6 +18,8 @@ from .api import (
|
|
|
15
18
|
evaluate_feedback_then_suggest_next_solution,
|
|
16
19
|
get_optimization_session_status,
|
|
17
20
|
handle_api_error,
|
|
21
|
+
send_heartbeat,
|
|
22
|
+
report_termination,
|
|
18
23
|
)
|
|
19
24
|
|
|
20
25
|
from . import __base_url__
|
|
@@ -37,11 +42,73 @@ from .utils import (
|
|
|
37
42
|
run_evaluation,
|
|
38
43
|
smooth_update,
|
|
39
44
|
format_number,
|
|
45
|
+
check_for_cli_updates,
|
|
40
46
|
)
|
|
41
47
|
|
|
42
48
|
install(show_locals=True)
|
|
43
49
|
console = Console()
|
|
44
50
|
|
|
51
|
+
# --- Global variable for heartbeat thread ---
|
|
52
|
+
heartbeat_thread = None
|
|
53
|
+
stop_heartbeat_event = threading.Event()
|
|
54
|
+
current_session_id_for_heartbeat = None
|
|
55
|
+
current_auth_headers_for_heartbeat = {}
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
# --- Heartbeat Sender Class ---
|
|
59
|
+
class HeartbeatSender(threading.Thread):
|
|
60
|
+
def __init__(self, session_id: str, auth_headers: dict, stop_event: threading.Event, interval: int = 30):
|
|
61
|
+
super().__init__(daemon=True) # Daemon thread exits when main thread exits
|
|
62
|
+
self.session_id = session_id
|
|
63
|
+
self.auth_headers = auth_headers
|
|
64
|
+
self.interval = interval
|
|
65
|
+
self.stop_event = stop_event
|
|
66
|
+
|
|
67
|
+
def run(self):
|
|
68
|
+
try:
|
|
69
|
+
while not self.stop_event.is_set():
|
|
70
|
+
if not send_heartbeat(self.session_id, self.auth_headers):
|
|
71
|
+
# send_heartbeat itself prints errors to stderr if it returns False
|
|
72
|
+
# No explicit HeartbeatSender log needed here unless more detail is desired for a False return
|
|
73
|
+
pass # Continue trying as per original logic
|
|
74
|
+
|
|
75
|
+
if self.stop_event.is_set(): # Check before waiting for responsiveness
|
|
76
|
+
break
|
|
77
|
+
|
|
78
|
+
self.stop_event.wait(self.interval) # Wait for interval or stop signal
|
|
79
|
+
|
|
80
|
+
except Exception as e:
|
|
81
|
+
# Catch any unexpected error in the loop to prevent silent thread death
|
|
82
|
+
print(
|
|
83
|
+
f"[ERROR HeartbeatSender] Unhandled exception in run loop for session {self.session_id}: {e}", file=sys.stderr
|
|
84
|
+
)
|
|
85
|
+
traceback.print_exc(file=sys.stderr)
|
|
86
|
+
# The loop will break due to the exception, and thread will terminate via finally.
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
# --- Signal Handling ---
|
|
90
|
+
def signal_handler(signum, frame):
|
|
91
|
+
signal_name = signal.Signals(signum).name
|
|
92
|
+
console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]")
|
|
93
|
+
|
|
94
|
+
# Stop heartbeat thread
|
|
95
|
+
stop_heartbeat_event.set()
|
|
96
|
+
if heartbeat_thread and heartbeat_thread.is_alive():
|
|
97
|
+
heartbeat_thread.join(timeout=2) # Give it a moment to stop
|
|
98
|
+
|
|
99
|
+
# Report termination (best effort)
|
|
100
|
+
if current_session_id_for_heartbeat:
|
|
101
|
+
report_termination(
|
|
102
|
+
session_id=current_session_id_for_heartbeat,
|
|
103
|
+
status_update="terminated",
|
|
104
|
+
reason=f"user_terminated_{signal_name.lower()}",
|
|
105
|
+
details=f"Process terminated by signal {signal_name} ({signum}).",
|
|
106
|
+
auth_headers=current_auth_headers_for_heartbeat,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
# Exit gracefully
|
|
110
|
+
sys.exit(0)
|
|
111
|
+
|
|
45
112
|
|
|
46
113
|
def perform_login(console: Console):
|
|
47
114
|
"""Handles the device login flow."""
|
|
@@ -161,6 +228,16 @@ def perform_login(console: Console):
|
|
|
161
228
|
|
|
162
229
|
def main() -> None:
|
|
163
230
|
"""Main function for the Weco CLI."""
|
|
231
|
+
# Setup signal handlers
|
|
232
|
+
signal.signal(signal.SIGINT, signal_handler)
|
|
233
|
+
signal.signal(signal.SIGTERM, signal_handler)
|
|
234
|
+
|
|
235
|
+
# --- Perform Update Check ---
|
|
236
|
+
# Import __pkg_version__ here to avoid circular import issues if it's also used in modules imported by cli.py
|
|
237
|
+
from . import __pkg_version__
|
|
238
|
+
|
|
239
|
+
check_for_cli_updates(__pkg_version__) # Call the imported function
|
|
240
|
+
|
|
164
241
|
# --- Argument Parsing ---
|
|
165
242
|
parser = argparse.ArgumentParser(
|
|
166
243
|
description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
|
|
@@ -194,11 +271,6 @@ def main() -> None:
|
|
|
194
271
|
type=str,
|
|
195
272
|
help="Description of additional instruction or path to a file containing additional instructions",
|
|
196
273
|
)
|
|
197
|
-
run_parser.add_argument(
|
|
198
|
-
"--preserve-source",
|
|
199
|
-
action="store_true",
|
|
200
|
-
help="If set, do not overwrite the original source file; only save modified versions in the runs directory",
|
|
201
|
-
)
|
|
202
274
|
|
|
203
275
|
# --- Logout Command ---
|
|
204
276
|
_ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
|
|
@@ -212,6 +284,10 @@ def main() -> None:
|
|
|
212
284
|
|
|
213
285
|
# --- Handle Run Command ---
|
|
214
286
|
elif args.command == "run":
|
|
287
|
+
global heartbeat_thread, current_session_id_for_heartbeat, current_auth_headers_for_heartbeat # Allow modification of globals
|
|
288
|
+
|
|
289
|
+
session_id = None # Initialize session_id
|
|
290
|
+
optimization_completed_normally = False # Flag for finally block
|
|
215
291
|
# --- Check Authentication ---
|
|
216
292
|
weco_api_key = load_weco_api_key()
|
|
217
293
|
llm_api_keys = read_api_keys_from_env() # Read keys from client environment
|
|
@@ -243,10 +319,9 @@ def main() -> None:
|
|
|
243
319
|
|
|
244
320
|
# --- Prepare API Call Arguments ---
|
|
245
321
|
auth_headers = {}
|
|
246
|
-
|
|
247
322
|
if weco_api_key:
|
|
248
323
|
auth_headers["Authorization"] = f"Bearer {weco_api_key}"
|
|
249
|
-
|
|
324
|
+
current_auth_headers_for_heartbeat = auth_headers # Store for signal handler
|
|
250
325
|
|
|
251
326
|
# --- Main Run Logic ---
|
|
252
327
|
try:
|
|
@@ -294,25 +369,30 @@ def main() -> None:
|
|
|
294
369
|
evaluator_config=evaluator_config,
|
|
295
370
|
search_policy_config=search_policy_config,
|
|
296
371
|
additional_instructions=additional_instructions,
|
|
297
|
-
api_keys=llm_api_keys,
|
|
298
|
-
auth_headers=auth_headers,
|
|
372
|
+
api_keys=llm_api_keys,
|
|
373
|
+
auth_headers=auth_headers,
|
|
299
374
|
timeout=timeout,
|
|
300
375
|
)
|
|
376
|
+
session_id = session_response["session_id"]
|
|
377
|
+
current_session_id_for_heartbeat = session_id # Store for signal handler/finally
|
|
378
|
+
|
|
379
|
+
# --- Start Heartbeat Thread ---
|
|
380
|
+
stop_heartbeat_event.clear() # Ensure event is clear before starting
|
|
381
|
+
heartbeat_thread = HeartbeatSender(session_id, auth_headers, stop_heartbeat_event)
|
|
382
|
+
heartbeat_thread.start()
|
|
301
383
|
|
|
302
384
|
# --- Live Update Loop ---
|
|
303
385
|
refresh_rate = 4
|
|
304
386
|
with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
|
|
305
387
|
# Define the runs directory (.runs/<session-id>)
|
|
306
|
-
session_id = session_response["session_id"]
|
|
307
388
|
runs_dir = pathlib.Path(args.log_dir) / session_id
|
|
308
389
|
runs_dir.mkdir(parents=True, exist_ok=True)
|
|
309
390
|
|
|
310
391
|
# Write the initial code string to the logs
|
|
311
392
|
write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=session_response["code"])
|
|
312
393
|
|
|
313
|
-
# Write the initial code string to the source file path
|
|
314
|
-
|
|
315
|
-
write_to_path(fp=source_fp, content=session_response["code"])
|
|
394
|
+
# Write the initial code string to the source file path
|
|
395
|
+
write_to_path(fp=source_fp, content=session_response["code"])
|
|
316
396
|
|
|
317
397
|
# Update the panels with the initial solution
|
|
318
398
|
summary_panel.set_session_id(session_id=session_id) # Add session id now that we have it
|
|
@@ -364,6 +444,9 @@ def main() -> None:
|
|
|
364
444
|
transition_delay=0.1,
|
|
365
445
|
)
|
|
366
446
|
|
|
447
|
+
# # Send initial heartbeat immediately after starting
|
|
448
|
+
# send_heartbeat(session_id, auth_headers)
|
|
449
|
+
|
|
367
450
|
# Run evaluation on the initial solution
|
|
368
451
|
term_out = run_evaluation(eval_command=args.eval_command)
|
|
369
452
|
|
|
@@ -392,14 +475,14 @@ def main() -> None:
|
|
|
392
475
|
auth_headers=auth_headers, # Pass Weco key if logged in
|
|
393
476
|
timeout=timeout,
|
|
394
477
|
)
|
|
478
|
+
|
|
395
479
|
# Save next solution (.runs/<session-id>/step_<step>.<extension>)
|
|
396
480
|
write_to_path(
|
|
397
481
|
fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"]
|
|
398
482
|
)
|
|
399
483
|
|
|
400
484
|
# Write the next solution to the source file
|
|
401
|
-
|
|
402
|
-
write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
|
|
485
|
+
write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
|
|
403
486
|
|
|
404
487
|
# Get the optimization session status for
|
|
405
488
|
# the best solution, its score, and the history to plot the tree
|
|
@@ -483,7 +566,7 @@ def main() -> None:
|
|
|
483
566
|
additional_instructions=args.additional_instructions
|
|
484
567
|
)
|
|
485
568
|
|
|
486
|
-
#
|
|
569
|
+
# Final evaluation report
|
|
487
570
|
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
|
|
488
571
|
session_id=session_id,
|
|
489
572
|
execution_output=term_out,
|
|
@@ -560,17 +643,76 @@ def main() -> None:
|
|
|
560
643
|
write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
|
|
561
644
|
|
|
562
645
|
# write the best solution to the source file
|
|
563
|
-
|
|
564
|
-
|
|
646
|
+
write_to_path(fp=source_fp, content=best_solution_content)
|
|
647
|
+
|
|
648
|
+
# Mark as completed normally for the finally block
|
|
649
|
+
optimization_completed_normally = True
|
|
565
650
|
|
|
566
651
|
console.print(end_optimization_layout)
|
|
567
652
|
|
|
568
653
|
except Exception as e:
|
|
654
|
+
# Catch errors during the main optimization loop or setup
|
|
569
655
|
try:
|
|
570
|
-
error_message = e.response.json()["detail"]
|
|
656
|
+
error_message = e.response.json()["detail"] # Try to get API error detail
|
|
571
657
|
except Exception:
|
|
572
|
-
error_message = str(e)
|
|
573
|
-
console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Error", border_style="red"))
|
|
574
|
-
# Print traceback for debugging
|
|
658
|
+
error_message = str(e) # Otherwise, use the exception string
|
|
659
|
+
console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Optimization Error", border_style="red"))
|
|
660
|
+
# Print traceback for debugging if needed (can be noisy)
|
|
575
661
|
# console.print_exception(show_locals=False)
|
|
576
|
-
|
|
662
|
+
|
|
663
|
+
# Ensure optimization_completed_normally is False
|
|
664
|
+
optimization_completed_normally = False
|
|
665
|
+
|
|
666
|
+
# Prepare details for termination report
|
|
667
|
+
error_details = traceback.format_exc()
|
|
668
|
+
|
|
669
|
+
# Exit code will be handled by finally block or sys.exit below
|
|
670
|
+
exit_code = 1 # Indicate error
|
|
671
|
+
# No sys.exit here, let finally block run
|
|
672
|
+
|
|
673
|
+
finally:
|
|
674
|
+
# This block runs whether the try block completed normally or raised an exception
|
|
675
|
+
|
|
676
|
+
# Stop heartbeat thread
|
|
677
|
+
stop_heartbeat_event.set()
|
|
678
|
+
if heartbeat_thread and heartbeat_thread.is_alive():
|
|
679
|
+
heartbeat_thread.join(timeout=2) # Give it a moment to stop
|
|
680
|
+
|
|
681
|
+
# Report final status if a session was started
|
|
682
|
+
if session_id:
|
|
683
|
+
final_status = "unknown"
|
|
684
|
+
final_reason = "unknown_termination"
|
|
685
|
+
final_details = None
|
|
686
|
+
|
|
687
|
+
if optimization_completed_normally:
|
|
688
|
+
final_status = "completed"
|
|
689
|
+
final_reason = "completed_successfully"
|
|
690
|
+
else:
|
|
691
|
+
# If an exception was caught and we have details
|
|
692
|
+
if "error_details" in locals():
|
|
693
|
+
final_status = "error"
|
|
694
|
+
final_reason = "error_cli_internal"
|
|
695
|
+
final_details = error_details
|
|
696
|
+
# else: # Should have been handled by signal handler if terminated by user
|
|
697
|
+
# Keep default 'unknown' if we somehow end up here without error/completion/signal
|
|
698
|
+
|
|
699
|
+
# Avoid reporting if terminated by signal handler (already reported)
|
|
700
|
+
# Check a flag or rely on status not being 'unknown'
|
|
701
|
+
if final_status != "unknown":
|
|
702
|
+
report_termination(
|
|
703
|
+
session_id=session_id,
|
|
704
|
+
status_update=final_status,
|
|
705
|
+
reason=final_reason,
|
|
706
|
+
details=final_details,
|
|
707
|
+
auth_headers=auth_headers,
|
|
708
|
+
)
|
|
709
|
+
|
|
710
|
+
# Ensure proper exit code if an error occurred
|
|
711
|
+
if not optimization_completed_normally and "exit_code" in locals() and exit_code != 0:
|
|
712
|
+
sys.exit(exit_code)
|
|
713
|
+
elif not optimization_completed_normally:
|
|
714
|
+
# Generic error exit if no specific code was set but try block failed
|
|
715
|
+
sys.exit(1)
|
|
716
|
+
else:
|
|
717
|
+
# Normal exit
|
|
718
|
+
sys.exit(0)
|
weco/panels.py
CHANGED
|
@@ -253,7 +253,7 @@ class MetricTreePanel:
|
|
|
253
253
|
# Make sure the metric tree is built before calling build_rich_tree
|
|
254
254
|
return Panel(
|
|
255
255
|
self._build_rich_tree(),
|
|
256
|
-
title="[bold]🔎 Exploring Solutions..." if not is_done else "[bold]🔎 Optimization Complete!",
|
|
256
|
+
title=("[bold]🔎 Exploring Solutions..." if not is_done else "[bold]🔎 Optimization Complete!"),
|
|
257
257
|
border_style="green",
|
|
258
258
|
expand=True,
|
|
259
259
|
padding=(0, 1),
|
weco/utils.py
CHANGED
|
@@ -7,6 +7,8 @@ from rich.layout import Layout
|
|
|
7
7
|
from rich.live import Live
|
|
8
8
|
from rich.panel import Panel
|
|
9
9
|
import pathlib
|
|
10
|
+
import requests
|
|
11
|
+
from packaging.version import parse as parse_version
|
|
10
12
|
|
|
11
13
|
|
|
12
14
|
# Env/arg helper functions
|
|
@@ -109,3 +111,33 @@ def run_evaluation(eval_command: str) -> str:
|
|
|
109
111
|
output += "\n"
|
|
110
112
|
output += result.stdout
|
|
111
113
|
return output
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
# Update Check Function
|
|
117
|
+
def check_for_cli_updates(current_version_str: str):
|
|
118
|
+
"""Checks PyPI for a newer version of the weco package and notifies the user."""
|
|
119
|
+
try:
|
|
120
|
+
pypi_url = "https://pypi.org/pypi/weco/json"
|
|
121
|
+
response = requests.get(pypi_url, timeout=5) # Short timeout for non-critical check
|
|
122
|
+
response.raise_for_status()
|
|
123
|
+
latest_version_str = response.json()["info"]["version"]
|
|
124
|
+
|
|
125
|
+
current_version = parse_version(current_version_str)
|
|
126
|
+
latest_version = parse_version(latest_version_str)
|
|
127
|
+
|
|
128
|
+
if latest_version > current_version:
|
|
129
|
+
yellow_start = "\033[93m"
|
|
130
|
+
reset_color = "\033[0m"
|
|
131
|
+
message = f"WARNING: New weco version ({latest_version_str}) available (you have {current_version_str}). Run: pip install --upgrade weco"
|
|
132
|
+
print(f"{yellow_start}{message}{reset_color}")
|
|
133
|
+
time.sleep(2) # Wait for 2 second
|
|
134
|
+
|
|
135
|
+
except requests.exceptions.RequestException:
|
|
136
|
+
# Silently fail on network errors, etc. Don't disrupt user.
|
|
137
|
+
pass
|
|
138
|
+
except (KeyError, ValueError):
|
|
139
|
+
# Handle cases where the PyPI response format might be unexpected
|
|
140
|
+
pass
|
|
141
|
+
except Exception:
|
|
142
|
+
# Catch any other unexpected error during the check
|
|
143
|
+
pass
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: weco
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.18
|
|
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
|
|
@@ -14,25 +14,33 @@ Description-Content-Type: text/markdown
|
|
|
14
14
|
License-File: LICENSE
|
|
15
15
|
Requires-Dist: requests
|
|
16
16
|
Requires-Dist: rich
|
|
17
|
+
Requires-Dist: packaging
|
|
17
18
|
Provides-Extra: dev
|
|
18
19
|
Requires-Dist: ruff; extra == "dev"
|
|
19
20
|
Requires-Dist: build; extra == "dev"
|
|
20
21
|
Requires-Dist: setuptools_scm; extra == "dev"
|
|
21
22
|
Dynamic: license-file
|
|
22
23
|
|
|
23
|
-
|
|
24
|
+
<div align="center">
|
|
24
25
|
|
|
25
|
-
|
|
26
|
+
# Weco: The Platform for Self-Improving Code
|
|
27
|
+
|
|
28
|
+
[](https://www.python.org)
|
|
29
|
+
[](https://docs.weco.ai/)
|
|
26
30
|
[](https://badge.fury.io/py/weco)
|
|
27
31
|
[](https://arxiv.org/abs/2502.13138)
|
|
28
32
|
|
|
33
|
+
</div>
|
|
34
|
+
|
|
35
|
+
---
|
|
36
|
+
|
|
29
37
|
Weco systematically optimizes your code, guided directly by your evaluation metrics.
|
|
30
38
|
|
|
31
39
|
Example applications include:
|
|
32
40
|
|
|
33
|
-
- **GPU Kernel Optimization**: Reimplement PyTorch functions using CUDA or Triton optimizing for `latency`, `throughput`, or `memory_bandwidth`.
|
|
34
|
-
- **Model Development**: Tune feature transformations or
|
|
35
|
-
- **Prompt Engineering**: Refine prompts for LLMs, optimizing for `win_rate`, `relevance`, or `format_adherence`
|
|
41
|
+
- **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`.
|
|
42
|
+
- **Model Development**: Tune feature transformations, architectures or [the whole training pipeline](/examples/spaceship-titanic/README.md), optimizing for `validation_accuracy`, `AUC`, or `Sharpe Ratio`.
|
|
43
|
+
- **Prompt Engineering**: Refine prompts for LLMs (e.g., for [math problems](/examples/prompt/README.md)), optimizing for `win_rate`, `relevance`, or `format_adherence`
|
|
36
44
|
|
|
37
45
|

|
|
38
46
|
|
|
@@ -62,29 +70,9 @@ The `weco` CLI leverages a tree search approach guided by Large Language Models
|
|
|
62
70
|
- **Anthropic:** `export ANTHROPIC_API_KEY="your_key_here"`
|
|
63
71
|
- **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
72
|
|
|
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.
|
|
79
|
-
|
|
80
73
|
---
|
|
81
74
|
|
|
82
|
-
##
|
|
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.
|
|
75
|
+
## Get Started
|
|
88
76
|
|
|
89
77
|
<div style="background-color: #fff3cd; border: 1px solid #ffeeba; padding: 15px; border-radius: 4px; margin-bottom: 15px;">
|
|
90
78
|
<strong>⚠️ Warning: Code Modification</strong><br>
|
|
@@ -93,15 +81,11 @@ The CLI has two main commands:
|
|
|
93
81
|
|
|
94
82
|
---
|
|
95
83
|
|
|
96
|
-
### `weco run` Command
|
|
97
|
-
|
|
98
|
-
This command starts the optimization process.
|
|
99
|
-
|
|
100
84
|
**Example: Optimizing Simple PyTorch Operations**
|
|
101
85
|
|
|
102
86
|
This basic example shows how to optimize a simple PyTorch function for speedup.
|
|
103
87
|
|
|
104
|
-
For more advanced examples, including [Triton](/examples/triton/README.md), [CUDA kernel optimization](/examples/cuda/README.md)
|
|
88
|
+
For more advanced examples, including [Triton](/examples/triton/README.md), [CUDA kernel optimization](/examples/cuda/README.md), [ML model optimization](/examples/spaceship-titanic/README.md), and [prompt engineering for math problems](https://github.com/WecoAI/weco-cli/tree/main/examples/prompt), please see the `README.md` files within the corresponding subdirectories under the [`examples/`](./examples/) folder.
|
|
105
89
|
|
|
106
90
|
```bash
|
|
107
91
|
# Navigate to the example directory
|
|
@@ -136,17 +120,12 @@ weco run --source optimize.py \
|
|
|
136
120
|
| `--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
121
|
| `--additional-instructions` | (Optional) Natural language description of specific instructions OR path to a file containing detailed instructions to guide the LLM. | No |
|
|
138
122
|
| `--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
123
|
|
|
141
124
|
---
|
|
142
125
|
|
|
143
|
-
###
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
```bash
|
|
148
|
-
weco logout
|
|
149
|
-
```
|
|
126
|
+
### Weco Dashboard
|
|
127
|
+
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
|
|
128
|
+

|
|
150
129
|
|
|
151
130
|
---
|
|
152
131
|
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
weco/__init__.py,sha256=npWmRgLxfVK69GdyxIujnI87xqmPCBrZWxxAxL_QQOc,478
|
|
2
|
+
weco/api.py,sha256=lJJ0j0-bABiQXDlRb43fCo7ky0N_HwfZgFdMktRKQ90,6635
|
|
3
|
+
weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
|
|
4
|
+
weco/cli.py,sha256=hrbZ24eoIHJdFzHoJ51WNGfdOzYtgskUvaor7JbMNWc,34973
|
|
5
|
+
weco/panels.py,sha256=pM_YGnmcXM_1CBcxo_EAzOV3g_4NFdLS4MqDqx7THbA,13563
|
|
6
|
+
weco/utils.py,sha256=LVTBo3dduJmhlbotcYoUW2nLx6IRtKs4eDFR52Qltcg,5244
|
|
7
|
+
weco-0.2.18.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
|
|
8
|
+
weco-0.2.18.dist-info/METADATA,sha256=0GcQBeAQrtRjaOUN0OmKGPvyJwD3y8eri6nY_He-5eY,9895
|
|
9
|
+
weco-0.2.18.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
|
|
10
|
+
weco-0.2.18.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
|
|
11
|
+
weco-0.2.18.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
|
|
12
|
+
weco-0.2.18.dist-info/RECORD,,
|
weco-0.2.16.dist-info/RECORD
DELETED
|
@@ -1,12 +0,0 @@
|
|
|
1
|
-
weco/__init__.py,sha256=HyOG-766A8if3MM5YvAM6CMJobEib3oM54UAJMk9GFA,426
|
|
2
|
-
weco/api.py,sha256=0OU1hEhN7sbIZ1zj8TeeB0tMaxXk6n6qw82FmcdK0ec,3111
|
|
3
|
-
weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
|
|
4
|
-
weco/cli.py,sha256=DPObmIes3tRAVswBGYOx5G_8zIjeIJFufOGnFt__XwE,28869
|
|
5
|
-
weco/panels.py,sha256=8DoTQC-epGpGjn-xDBcqelC5BKaX7JXnrJ97LInEbRU,13561
|
|
6
|
-
weco/utils.py,sha256=hhIebUPnetFMfNSFfcsKVw1TSpeu_Zw3rBPPnxDie0U,3911
|
|
7
|
-
weco-0.2.16.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
|
|
8
|
-
weco-0.2.16.dist-info/METADATA,sha256=CFe6w8pRnMh-zDuyebtqIOFfIuf6JEJFqgVdUPSTO6s,10749
|
|
9
|
-
weco-0.2.16.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
|
|
10
|
-
weco-0.2.16.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
|
|
11
|
-
weco-0.2.16.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
|
|
12
|
-
weco-0.2.16.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|