weco 0.2.18__py3-none-any.whl → 0.2.20__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/api.py +24 -40
- weco/cli.py +243 -247
- weco/panels.py +13 -10
- {weco-0.2.18.dist-info → weco-0.2.20.dist-info}/METADATA +26 -15
- weco-0.2.20.dist-info/RECORD +12 -0
- {weco-0.2.18.dist-info → weco-0.2.20.dist-info}/WHEEL +1 -1
- weco-0.2.18.dist-info/RECORD +0 -12
- {weco-0.2.18.dist-info → weco-0.2.20.dist-info}/entry_points.txt +0 -0
- {weco-0.2.18.dist-info → weco-0.2.20.dist-info}/licenses/LICENSE +0 -0
- {weco-0.2.18.dist-info → weco-0.2.20.dist-info}/top_level.txt +0 -0
weco/api.py
CHANGED
|
@@ -17,7 +17,7 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Con
|
|
|
17
17
|
# sys.exit(1)
|
|
18
18
|
|
|
19
19
|
|
|
20
|
-
def
|
|
20
|
+
def start_optimization_run(
|
|
21
21
|
console: rich.console.Console,
|
|
22
22
|
source_code: str,
|
|
23
23
|
evaluation_command: str,
|
|
@@ -29,14 +29,14 @@ def start_optimization_session(
|
|
|
29
29
|
search_policy_config: Dict[str, Any],
|
|
30
30
|
additional_instructions: str = None,
|
|
31
31
|
api_keys: Dict[str, Any] = {},
|
|
32
|
-
auth_headers: dict = {},
|
|
32
|
+
auth_headers: dict = {},
|
|
33
33
|
timeout: int = 800,
|
|
34
34
|
) -> Dict[str, Any]:
|
|
35
|
-
"""Start the optimization
|
|
35
|
+
"""Start the optimization run."""
|
|
36
36
|
with console.status("[bold green]Starting Optimization..."):
|
|
37
37
|
try:
|
|
38
38
|
response = requests.post(
|
|
39
|
-
f"{__base_url__}/
|
|
39
|
+
f"{__base_url__}/runs",
|
|
40
40
|
json={
|
|
41
41
|
"source_code": source_code,
|
|
42
42
|
"additional_instructions": additional_instructions,
|
|
@@ -49,37 +49,37 @@ def start_optimization_session(
|
|
|
49
49
|
},
|
|
50
50
|
"metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
|
|
51
51
|
},
|
|
52
|
-
headers=auth_headers,
|
|
52
|
+
headers=auth_headers,
|
|
53
53
|
timeout=timeout,
|
|
54
54
|
)
|
|
55
55
|
response.raise_for_status()
|
|
56
56
|
return response.json()
|
|
57
57
|
except requests.exceptions.HTTPError as e:
|
|
58
58
|
handle_api_error(e, console)
|
|
59
|
-
sys.exit(1)
|
|
59
|
+
sys.exit(1)
|
|
60
60
|
except requests.exceptions.RequestException as e:
|
|
61
|
-
console.print(f"[bold red]Network Error starting
|
|
61
|
+
console.print(f"[bold red]Network Error starting run: {e}[/]")
|
|
62
62
|
sys.exit(1)
|
|
63
63
|
|
|
64
64
|
|
|
65
65
|
def evaluate_feedback_then_suggest_next_solution(
|
|
66
|
-
|
|
66
|
+
run_id: str,
|
|
67
67
|
execution_output: str,
|
|
68
68
|
additional_instructions: str = None,
|
|
69
69
|
api_keys: Dict[str, Any] = {},
|
|
70
|
-
auth_headers: dict = {},
|
|
70
|
+
auth_headers: dict = {},
|
|
71
71
|
timeout: int = 800,
|
|
72
72
|
) -> Dict[str, Any]:
|
|
73
73
|
"""Evaluate the feedback and suggest the next solution."""
|
|
74
74
|
try:
|
|
75
75
|
response = requests.post(
|
|
76
|
-
f"{__base_url__}/
|
|
76
|
+
f"{__base_url__}/runs/{run_id}/suggest",
|
|
77
77
|
json={
|
|
78
78
|
"execution_output": execution_output,
|
|
79
79
|
"additional_instructions": additional_instructions,
|
|
80
80
|
"metadata": {**api_keys},
|
|
81
81
|
},
|
|
82
|
-
headers=auth_headers,
|
|
82
|
+
headers=auth_headers,
|
|
83
83
|
timeout=timeout,
|
|
84
84
|
)
|
|
85
85
|
response.raise_for_status()
|
|
@@ -93,16 +93,13 @@ def evaluate_feedback_then_suggest_next_solution(
|
|
|
93
93
|
raise # Re-raise the exception
|
|
94
94
|
|
|
95
95
|
|
|
96
|
-
def
|
|
97
|
-
|
|
96
|
+
def get_optimization_run_status(
|
|
97
|
+
run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
|
|
98
98
|
) -> Dict[str, Any]:
|
|
99
|
-
"""Get the current status of the optimization
|
|
99
|
+
"""Get the current status of the optimization run."""
|
|
100
100
|
try:
|
|
101
101
|
response = requests.get(
|
|
102
|
-
f"{__base_url__}/
|
|
103
|
-
params={"include_history": include_history},
|
|
104
|
-
headers=auth_headers,
|
|
105
|
-
timeout=timeout,
|
|
102
|
+
f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
|
|
106
103
|
)
|
|
107
104
|
response.raise_for_status()
|
|
108
105
|
return response.json()
|
|
@@ -114,42 +111,30 @@ def get_optimization_session_status(
|
|
|
114
111
|
raise # Re-raise
|
|
115
112
|
|
|
116
113
|
|
|
117
|
-
def send_heartbeat(
|
|
118
|
-
session_id: str,
|
|
119
|
-
auth_headers: dict = {},
|
|
120
|
-
timeout: int = 10, # Shorter timeout for non-critical heartbeat
|
|
121
|
-
) -> bool:
|
|
114
|
+
def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: int = 10) -> bool:
|
|
122
115
|
"""Send a heartbeat signal to the backend."""
|
|
123
116
|
try:
|
|
124
|
-
response = requests.put(f"{__base_url__}/
|
|
125
|
-
response.raise_for_status()
|
|
117
|
+
response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
|
|
118
|
+
response.raise_for_status()
|
|
126
119
|
return True
|
|
127
120
|
except requests.exceptions.HTTPError as e:
|
|
128
|
-
# Log non-critical errors like 409 Conflict (session not running)
|
|
129
121
|
if e.response.status_code == 409:
|
|
130
|
-
print(f"Heartbeat ignored:
|
|
122
|
+
print(f"Heartbeat ignored: Run {run_id} is not running.", file=sys.stderr)
|
|
131
123
|
else:
|
|
132
|
-
print(f"Heartbeat failed for
|
|
133
|
-
# Don't exit, just report failure
|
|
124
|
+
print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
|
|
134
125
|
return False
|
|
135
126
|
except requests.exceptions.RequestException as e:
|
|
136
|
-
|
|
137
|
-
print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
|
|
127
|
+
print(f"Heartbeat network error for run {run_id}: {e}", file=sys.stderr)
|
|
138
128
|
return False
|
|
139
129
|
|
|
140
130
|
|
|
141
131
|
def report_termination(
|
|
142
|
-
|
|
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
|
|
132
|
+
run_id: str, status_update: str, reason: str, details: Optional[str] = None, auth_headers: dict = {}, timeout: int = 30
|
|
148
133
|
) -> bool:
|
|
149
134
|
"""Report the termination reason to the backend."""
|
|
150
135
|
try:
|
|
151
136
|
response = requests.post(
|
|
152
|
-
f"{__base_url__}/
|
|
137
|
+
f"{__base_url__}/runs/{run_id}/terminate",
|
|
153
138
|
json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
|
|
154
139
|
headers=auth_headers,
|
|
155
140
|
timeout=timeout,
|
|
@@ -157,6 +142,5 @@ def report_termination(
|
|
|
157
142
|
response.raise_for_status()
|
|
158
143
|
return True
|
|
159
144
|
except requests.exceptions.RequestException as e:
|
|
160
|
-
|
|
161
|
-
print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
|
|
145
|
+
print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
|
|
162
146
|
return False
|
weco/cli.py
CHANGED
|
@@ -14,9 +14,9 @@ from rich.panel import Panel
|
|
|
14
14
|
from rich.traceback import install
|
|
15
15
|
from rich.prompt import Prompt
|
|
16
16
|
from .api import (
|
|
17
|
-
|
|
17
|
+
start_optimization_run,
|
|
18
18
|
evaluate_feedback_then_suggest_next_solution,
|
|
19
|
-
|
|
19
|
+
get_optimization_run_status,
|
|
20
20
|
handle_api_error,
|
|
21
21
|
send_heartbeat,
|
|
22
22
|
report_termination,
|
|
@@ -51,15 +51,15 @@ console = Console()
|
|
|
51
51
|
# --- Global variable for heartbeat thread ---
|
|
52
52
|
heartbeat_thread = None
|
|
53
53
|
stop_heartbeat_event = threading.Event()
|
|
54
|
-
|
|
54
|
+
current_run_id_for_heartbeat = None
|
|
55
55
|
current_auth_headers_for_heartbeat = {}
|
|
56
56
|
|
|
57
57
|
|
|
58
58
|
# --- Heartbeat Sender Class ---
|
|
59
59
|
class HeartbeatSender(threading.Thread):
|
|
60
|
-
def __init__(self,
|
|
60
|
+
def __init__(self, run_id: str, auth_headers: dict, stop_event: threading.Event, interval: int = 30):
|
|
61
61
|
super().__init__(daemon=True) # Daemon thread exits when main thread exits
|
|
62
|
-
self.
|
|
62
|
+
self.run_id = run_id
|
|
63
63
|
self.auth_headers = auth_headers
|
|
64
64
|
self.interval = interval
|
|
65
65
|
self.stop_event = stop_event
|
|
@@ -67,7 +67,7 @@ class HeartbeatSender(threading.Thread):
|
|
|
67
67
|
def run(self):
|
|
68
68
|
try:
|
|
69
69
|
while not self.stop_event.is_set():
|
|
70
|
-
if not send_heartbeat(self.
|
|
70
|
+
if not send_heartbeat(self.run_id, self.auth_headers):
|
|
71
71
|
# send_heartbeat itself prints errors to stderr if it returns False
|
|
72
72
|
# No explicit HeartbeatSender log needed here unless more detail is desired for a False return
|
|
73
73
|
pass # Continue trying as per original logic
|
|
@@ -79,9 +79,7 @@ class HeartbeatSender(threading.Thread):
|
|
|
79
79
|
|
|
80
80
|
except Exception as e:
|
|
81
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
|
-
)
|
|
82
|
+
print(f"[ERROR HeartbeatSender] Unhandled exception in run loop for run {self.run_id}: {e}", file=sys.stderr)
|
|
85
83
|
traceback.print_exc(file=sys.stderr)
|
|
86
84
|
# The loop will break due to the exception, and thread will terminate via finally.
|
|
87
85
|
|
|
@@ -97,13 +95,14 @@ def signal_handler(signum, frame):
|
|
|
97
95
|
heartbeat_thread.join(timeout=2) # Give it a moment to stop
|
|
98
96
|
|
|
99
97
|
# Report termination (best effort)
|
|
100
|
-
if
|
|
98
|
+
if current_run_id_for_heartbeat:
|
|
101
99
|
report_termination(
|
|
102
|
-
|
|
100
|
+
run_id=current_run_id_for_heartbeat,
|
|
103
101
|
status_update="terminated",
|
|
104
102
|
reason=f"user_terminated_{signal_name.lower()}",
|
|
105
103
|
details=f"Process terminated by signal {signal_name} ({signum}).",
|
|
106
104
|
auth_headers=current_auth_headers_for_heartbeat,
|
|
105
|
+
timeout=3,
|
|
107
106
|
)
|
|
108
107
|
|
|
109
108
|
# Exit gracefully
|
|
@@ -158,7 +157,7 @@ def perform_login(console: Console):
|
|
|
158
157
|
|
|
159
158
|
try:
|
|
160
159
|
token_response = requests.post(
|
|
161
|
-
f"{__base_url__}/auth/device/token",
|
|
160
|
+
f"{__base_url__}/auth/device/token",
|
|
162
161
|
json={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code},
|
|
163
162
|
)
|
|
164
163
|
|
|
@@ -172,12 +171,10 @@ def perform_login(console: Console):
|
|
|
172
171
|
# Unexpected 202 response format
|
|
173
172
|
console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
|
|
174
173
|
return False
|
|
175
|
-
|
|
176
174
|
# Check for standard OAuth2 errors (often 400 Bad Request)
|
|
177
175
|
elif token_response.status_code == 400:
|
|
178
176
|
token_data = token_response.json()
|
|
179
177
|
error_code = token_data.get("error", "unknown_error")
|
|
180
|
-
# NOTE: Removed "authorization_pending" check from here
|
|
181
178
|
if error_code == "slow_down":
|
|
182
179
|
interval += 5 # Increase polling interval if instructed
|
|
183
180
|
live_status.update(f"Waiting... (slowing down polling to {interval}s)")
|
|
@@ -195,7 +192,6 @@ def perform_login(console: Console):
|
|
|
195
192
|
|
|
196
193
|
# Check for other non-200/non-202/non-400 HTTP errors
|
|
197
194
|
token_response.raise_for_status()
|
|
198
|
-
|
|
199
195
|
# If successful (200 OK and no 'error' field)
|
|
200
196
|
token_data = token_response.json()
|
|
201
197
|
if "access_token" in token_data:
|
|
@@ -206,19 +202,17 @@ def perform_login(console: Console):
|
|
|
206
202
|
else:
|
|
207
203
|
# Unexpected successful response format
|
|
208
204
|
console.print("\n[bold red]Error:[/] Received unexpected response from server during polling.")
|
|
209
|
-
print(token_data)
|
|
205
|
+
print(token_data)
|
|
210
206
|
return False
|
|
211
|
-
|
|
212
207
|
except requests.exceptions.RequestException as e:
|
|
213
208
|
# Handle network errors during polling gracefully
|
|
214
209
|
live_status.update("Waiting... (network error, retrying)")
|
|
215
210
|
console.print(f"\n[bold yellow]Warning:[/] Network error during polling: {e}. Retrying...")
|
|
216
|
-
# Optional: implement backoff strategy
|
|
217
211
|
time.sleep(interval * 2) # Simple backoff
|
|
218
|
-
|
|
219
|
-
except requests.exceptions.HTTPError as e: # Catch HTTPError specifically for handle_api_error
|
|
212
|
+
except requests.exceptions.HTTPError as e:
|
|
220
213
|
handle_api_error(e, console)
|
|
221
|
-
except requests.exceptions.RequestException as e:
|
|
214
|
+
except requests.exceptions.RequestException as e:
|
|
215
|
+
# Catch other request errors
|
|
222
216
|
console.print(f"\n[bold red]Network Error:[/] {e}")
|
|
223
217
|
return False
|
|
224
218
|
except Exception as e:
|
|
@@ -233,72 +227,90 @@ def main() -> None:
|
|
|
233
227
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
234
228
|
|
|
235
229
|
# --- 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
230
|
from . import __pkg_version__
|
|
238
231
|
|
|
239
|
-
check_for_cli_updates(__pkg_version__)
|
|
232
|
+
check_for_cli_updates(__pkg_version__)
|
|
240
233
|
|
|
241
234
|
# --- Argument Parsing ---
|
|
242
235
|
parser = argparse.ArgumentParser(
|
|
243
236
|
description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
|
|
244
237
|
)
|
|
245
|
-
|
|
246
|
-
subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True) # Make command required
|
|
238
|
+
subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True)
|
|
247
239
|
|
|
248
240
|
# --- Run Command ---
|
|
249
241
|
run_parser = subparsers.add_parser(
|
|
250
|
-
"run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter
|
|
242
|
+
"run", help="Run code optimization", formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False
|
|
251
243
|
)
|
|
252
244
|
# Add arguments specific to the 'run' command to the run_parser
|
|
253
|
-
run_parser.add_argument("--source", type=str, required=True, help="Path to the source code file (e.g. optimize.py)")
|
|
254
245
|
run_parser.add_argument(
|
|
255
|
-
"
|
|
246
|
+
"-s",
|
|
247
|
+
"--source",
|
|
248
|
+
type=str,
|
|
249
|
+
required=True,
|
|
250
|
+
help="Path to the source code file that will be optimized (e.g., `optimize.py`)",
|
|
251
|
+
)
|
|
252
|
+
run_parser.add_argument(
|
|
253
|
+
"-c",
|
|
254
|
+
"--eval-command",
|
|
255
|
+
type=str,
|
|
256
|
+
required=True,
|
|
257
|
+
help="Command to run for evaluation (e.g. 'python eval.py --arg1=val1').",
|
|
258
|
+
)
|
|
259
|
+
run_parser.add_argument(
|
|
260
|
+
"-m",
|
|
261
|
+
"--metric",
|
|
262
|
+
type=str,
|
|
263
|
+
required=True,
|
|
264
|
+
help="Metric to optimize (e.g. 'accuracy', 'loss', 'f1_score') that is printed to the terminal by the eval command.",
|
|
256
265
|
)
|
|
257
|
-
run_parser.add_argument("--metric", type=str, required=True, help="Metric to optimize")
|
|
258
266
|
run_parser.add_argument(
|
|
259
|
-
"
|
|
267
|
+
"-g",
|
|
268
|
+
"--goal",
|
|
260
269
|
type=str,
|
|
261
|
-
choices=["
|
|
270
|
+
choices=["maximize", "max", "minimize", "min"],
|
|
262
271
|
required=True,
|
|
263
|
-
help="Specify '
|
|
272
|
+
help="Specify 'maximize'/'max' to maximize the metric or 'minimize'/'min' to minimize it.",
|
|
264
273
|
)
|
|
265
|
-
run_parser.add_argument("--steps", type=int,
|
|
266
|
-
run_parser.add_argument("--model", type=str, required=True, help="Model to use for optimization")
|
|
267
|
-
run_parser.add_argument("--log-dir", type=str, default=".runs", help="Directory to store logs and results")
|
|
274
|
+
run_parser.add_argument("-n", "--steps", type=int, default=100, help="Number of steps to run. Defaults to 100.")
|
|
268
275
|
run_parser.add_argument(
|
|
276
|
+
"-M",
|
|
277
|
+
"--model",
|
|
278
|
+
type=str,
|
|
279
|
+
default=None,
|
|
280
|
+
help="Model to use for optimization. Defaults to `o4-mini` when `OPENAI_API_KEY` is set, `claude-3-7-sonnet-20250219` when `ANTHROPIC_API_KEY` is set, and `gemini-2.5-pro-exp-03-25` when `GEMINI_API_KEY` is set. When multiple keys are set, the priority is `OPENAI_API_KEY` > `ANTHROPIC_API_KEY` > `GEMINI_API_KEY`.",
|
|
281
|
+
)
|
|
282
|
+
run_parser.add_argument(
|
|
283
|
+
"-l", "--log-dir", type=str, default=".runs", help="Directory to store logs and results. Defaults to `.runs`."
|
|
284
|
+
)
|
|
285
|
+
run_parser.add_argument(
|
|
286
|
+
"-i",
|
|
269
287
|
"--additional-instructions",
|
|
270
288
|
default=None,
|
|
271
289
|
type=str,
|
|
272
|
-
help="Description of additional instruction or path to a file containing additional instructions",
|
|
290
|
+
help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
|
|
273
291
|
)
|
|
274
292
|
|
|
275
|
-
# --- Logout Command ---
|
|
276
293
|
_ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
|
|
277
|
-
|
|
278
294
|
args = parser.parse_args()
|
|
279
295
|
|
|
280
|
-
# --- Handle Logout Command ---
|
|
281
296
|
if args.command == "logout":
|
|
282
297
|
clear_api_key()
|
|
283
298
|
sys.exit(0)
|
|
284
|
-
|
|
285
|
-
# --- Handle Run Command ---
|
|
286
299
|
elif args.command == "run":
|
|
287
|
-
global heartbeat_thread,
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
# --- Check Authentication ---
|
|
300
|
+
global heartbeat_thread, current_run_id_for_heartbeat, current_auth_headers_for_heartbeat # Allow modification of globals
|
|
301
|
+
run_id = None # Initialize run_id (we receive this from the API after starting the run)
|
|
302
|
+
optimization_completed_normally = False
|
|
303
|
+
user_stop_requested_flag = False
|
|
292
304
|
weco_api_key = load_weco_api_key()
|
|
293
305
|
llm_api_keys = read_api_keys_from_env() # Read keys from client environment
|
|
294
306
|
|
|
307
|
+
# --- Login/Authentication Handling ---
|
|
295
308
|
if not weco_api_key:
|
|
296
309
|
login_choice = Prompt.ask(
|
|
297
310
|
"Log in to Weco to save run history or use anonymously? ([bold]L[/]ogin / [bold]S[/]kip)",
|
|
298
311
|
choices=["l", "s"],
|
|
299
312
|
default="s",
|
|
300
313
|
).lower()
|
|
301
|
-
|
|
302
314
|
if login_choice == "l":
|
|
303
315
|
console.print("[cyan]Starting login process...[/]")
|
|
304
316
|
if not perform_login(console):
|
|
@@ -308,7 +320,6 @@ def main() -> None:
|
|
|
308
320
|
if not weco_api_key:
|
|
309
321
|
console.print("[bold red]Error: Login completed but failed to retrieve API key.[/]")
|
|
310
322
|
sys.exit(1)
|
|
311
|
-
|
|
312
323
|
elif login_choice == "s":
|
|
313
324
|
console.print("[yellow]Proceeding anonymously. LLM API keys must be provided via environment variables.[/]")
|
|
314
325
|
if not llm_api_keys:
|
|
@@ -317,7 +328,6 @@ def main() -> None:
|
|
|
317
328
|
)
|
|
318
329
|
sys.exit(1)
|
|
319
330
|
|
|
320
|
-
# --- Prepare API Call Arguments ---
|
|
321
331
|
auth_headers = {}
|
|
322
332
|
if weco_api_key:
|
|
323
333
|
auth_headers["Authorization"] = f"Bearer {weco_api_key}"
|
|
@@ -325,11 +335,23 @@ def main() -> None:
|
|
|
325
335
|
|
|
326
336
|
# --- Main Run Logic ---
|
|
327
337
|
try:
|
|
328
|
-
# ---
|
|
338
|
+
# --- Read Command Line Arguments ---
|
|
329
339
|
evaluation_command = args.eval_command
|
|
330
340
|
metric_name = args.metric
|
|
331
|
-
maximize = args.
|
|
341
|
+
maximize = args.goal in ["maximize", "max"]
|
|
332
342
|
steps = args.steps
|
|
343
|
+
# Determine the model to use
|
|
344
|
+
if args.model is None:
|
|
345
|
+
if "OPENAI_API_KEY" in llm_api_keys:
|
|
346
|
+
args.model = "o4-mini"
|
|
347
|
+
elif "ANTHROPIC_API_KEY" in llm_api_keys:
|
|
348
|
+
args.model = "claude-3-7-sonnet-20250219"
|
|
349
|
+
elif "GEMINI_API_KEY" in llm_api_keys:
|
|
350
|
+
args.model = "gemini-2.5-pro-exp-03-25"
|
|
351
|
+
else:
|
|
352
|
+
raise ValueError(
|
|
353
|
+
"No LLM API keys found in environment. Please set one of the following: OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY."
|
|
354
|
+
)
|
|
333
355
|
code_generator_config = {"model": args.model}
|
|
334
356
|
evaluator_config = {"model": args.model, "include_analysis": True}
|
|
335
357
|
search_policy_config = {
|
|
@@ -337,13 +359,9 @@ def main() -> None:
|
|
|
337
359
|
"debug_prob": 0.5,
|
|
338
360
|
"max_debug_depth": max(1, math.ceil(0.1 * steps)),
|
|
339
361
|
}
|
|
340
|
-
# API request timeout
|
|
341
362
|
timeout = 800
|
|
342
|
-
# Read additional instructions
|
|
343
363
|
additional_instructions = read_additional_instructions(additional_instructions=args.additional_instructions)
|
|
344
|
-
# Read source code path
|
|
345
364
|
source_fp = pathlib.Path(args.source)
|
|
346
|
-
# Read source code content
|
|
347
365
|
source_code = read_from_path(fp=source_fp, is_json=False)
|
|
348
366
|
|
|
349
367
|
# --- Panel Initialization ---
|
|
@@ -357,8 +375,8 @@ def main() -> None:
|
|
|
357
375
|
layout = create_optimization_layout()
|
|
358
376
|
end_optimization_layout = create_end_optimization_layout()
|
|
359
377
|
|
|
360
|
-
# --- Start Optimization
|
|
361
|
-
|
|
378
|
+
# --- Start Optimization Run ---
|
|
379
|
+
run_response = start_optimization_run(
|
|
362
380
|
console=console,
|
|
363
381
|
source_code=source_code,
|
|
364
382
|
evaluation_command=evaluation_command,
|
|
@@ -373,42 +391,39 @@ def main() -> None:
|
|
|
373
391
|
auth_headers=auth_headers,
|
|
374
392
|
timeout=timeout,
|
|
375
393
|
)
|
|
376
|
-
|
|
377
|
-
|
|
394
|
+
run_id = run_response["run_id"]
|
|
395
|
+
current_run_id_for_heartbeat = run_id
|
|
378
396
|
|
|
379
397
|
# --- Start Heartbeat Thread ---
|
|
380
|
-
stop_heartbeat_event.clear()
|
|
381
|
-
heartbeat_thread = HeartbeatSender(
|
|
398
|
+
stop_heartbeat_event.clear()
|
|
399
|
+
heartbeat_thread = HeartbeatSender(run_id, auth_headers, stop_heartbeat_event)
|
|
382
400
|
heartbeat_thread.start()
|
|
383
401
|
|
|
384
402
|
# --- Live Update Loop ---
|
|
385
403
|
refresh_rate = 4
|
|
386
404
|
with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
|
|
387
|
-
# Define the runs directory (.runs/<
|
|
388
|
-
runs_dir = pathlib.Path(args.log_dir) /
|
|
405
|
+
# Define the runs directory (.runs/<run-id>) to store logs and results
|
|
406
|
+
runs_dir = pathlib.Path(args.log_dir) / run_id
|
|
389
407
|
runs_dir.mkdir(parents=True, exist_ok=True)
|
|
390
|
-
|
|
391
408
|
# Write the initial code string to the logs
|
|
392
|
-
write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=
|
|
393
|
-
|
|
409
|
+
write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=run_response["code"])
|
|
394
410
|
# Write the initial code string to the source file path
|
|
395
|
-
write_to_path(fp=source_fp, content=
|
|
411
|
+
write_to_path(fp=source_fp, content=run_response["code"])
|
|
396
412
|
|
|
397
413
|
# Update the panels with the initial solution
|
|
398
|
-
summary_panel.
|
|
414
|
+
summary_panel.set_run_id(run_id=run_id) # Add run id now that we have it
|
|
399
415
|
# Set the step of the progress bar
|
|
400
416
|
summary_panel.set_step(step=0)
|
|
401
417
|
# Update the token counts
|
|
402
|
-
summary_panel.update_token_counts(usage=
|
|
403
|
-
|
|
404
|
-
plan_panel.update(plan=session_response["plan"])
|
|
418
|
+
summary_panel.update_token_counts(usage=run_response["usage"])
|
|
419
|
+
plan_panel.update(plan=run_response["plan"])
|
|
405
420
|
# Build the metric tree
|
|
406
421
|
tree_panel.build_metric_tree(
|
|
407
422
|
nodes=[
|
|
408
423
|
{
|
|
409
|
-
"solution_id":
|
|
424
|
+
"solution_id": run_response["solution_id"],
|
|
410
425
|
"parent_id": None,
|
|
411
|
-
"code":
|
|
426
|
+
"code": run_response["code"],
|
|
412
427
|
"step": 0,
|
|
413
428
|
"metric_value": None,
|
|
414
429
|
"is_buggy": False,
|
|
@@ -416,15 +431,11 @@ def main() -> None:
|
|
|
416
431
|
]
|
|
417
432
|
)
|
|
418
433
|
# Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
|
|
419
|
-
tree_panel.set_unevaluated_node(node_id=
|
|
434
|
+
tree_panel.set_unevaluated_node(node_id=run_response["solution_id"])
|
|
420
435
|
# Update the solution panels with the initial solution and get the panel displays
|
|
421
436
|
solution_panels.update(
|
|
422
437
|
current_node=Node(
|
|
423
|
-
id=
|
|
424
|
-
parent_id=None,
|
|
425
|
-
code=session_response["code"],
|
|
426
|
-
metric=None,
|
|
427
|
-
is_buggy=False,
|
|
438
|
+
id=run_response["solution_id"], parent_id=None, code=run_response["code"], metric=None, is_buggy=False
|
|
428
439
|
),
|
|
429
440
|
best_node=None,
|
|
430
441
|
)
|
|
@@ -444,12 +455,8 @@ def main() -> None:
|
|
|
444
455
|
transition_delay=0.1,
|
|
445
456
|
)
|
|
446
457
|
|
|
447
|
-
# # Send initial heartbeat immediately after starting
|
|
448
|
-
# send_heartbeat(session_id, auth_headers)
|
|
449
|
-
|
|
450
458
|
# Run evaluation on the initial solution
|
|
451
459
|
term_out = run_evaluation(eval_command=args.eval_command)
|
|
452
|
-
|
|
453
460
|
# Update the evaluation output panel
|
|
454
461
|
eval_output_panel.update(output=term_out)
|
|
455
462
|
smooth_update(
|
|
@@ -465,40 +472,50 @@ def main() -> None:
|
|
|
465
472
|
current_additional_instructions = read_additional_instructions(
|
|
466
473
|
additional_instructions=args.additional_instructions
|
|
467
474
|
)
|
|
475
|
+
if run_id:
|
|
476
|
+
try:
|
|
477
|
+
current_status_response = get_optimization_run_status(
|
|
478
|
+
run_id=run_id, include_history=False, timeout=30, auth_headers=auth_headers
|
|
479
|
+
)
|
|
480
|
+
current_run_status_val = current_status_response.get("status")
|
|
481
|
+
if current_run_status_val == "stopping":
|
|
482
|
+
console.print("\n[bold yellow]Stop request received. Terminating run gracefully...[/]")
|
|
483
|
+
user_stop_requested_flag = True
|
|
484
|
+
break
|
|
485
|
+
except requests.exceptions.RequestException as e:
|
|
486
|
+
console.print(
|
|
487
|
+
f"\n[bold red]Warning: Could not check run status: {e}. Continuing optimization...[/]"
|
|
488
|
+
)
|
|
489
|
+
except Exception as e:
|
|
490
|
+
console.print(
|
|
491
|
+
f"\n[bold red]Warning: Error checking run status: {e}. Continuing optimization...[/]"
|
|
492
|
+
)
|
|
468
493
|
|
|
469
494
|
# Send feedback and get next suggestion
|
|
470
495
|
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
|
|
471
|
-
|
|
496
|
+
run_id=run_id,
|
|
472
497
|
execution_output=term_out,
|
|
473
|
-
additional_instructions=current_additional_instructions,
|
|
474
|
-
api_keys=llm_api_keys,
|
|
475
|
-
auth_headers=auth_headers,
|
|
498
|
+
additional_instructions=current_additional_instructions,
|
|
499
|
+
api_keys=llm_api_keys,
|
|
500
|
+
auth_headers=auth_headers,
|
|
476
501
|
timeout=timeout,
|
|
477
502
|
)
|
|
478
|
-
|
|
479
|
-
# Save next solution (.runs/<session-id>/step_<step>.<extension>)
|
|
503
|
+
# Save next solution (.runs/<run-id>/step_<step>.<extension>)
|
|
480
504
|
write_to_path(
|
|
481
505
|
fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"]
|
|
482
506
|
)
|
|
483
|
-
|
|
484
507
|
# Write the next solution to the source file
|
|
485
508
|
write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
# the best solution, its score, and the history to plot the tree
|
|
489
|
-
status_response = get_optimization_session_status(
|
|
490
|
-
session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
|
|
509
|
+
status_response = get_optimization_run_status(
|
|
510
|
+
run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
|
|
491
511
|
)
|
|
492
|
-
|
|
493
|
-
# Update the step of the progress bar
|
|
512
|
+
# Update the step of the progress bar, token counts, plan and metric tree
|
|
494
513
|
summary_panel.set_step(step=step)
|
|
495
|
-
# Update the token counts
|
|
496
514
|
summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
|
|
497
|
-
# Update the plan
|
|
498
515
|
plan_panel.update(plan=eval_and_next_solution_response["plan"])
|
|
499
|
-
|
|
500
|
-
|
|
501
|
-
|
|
516
|
+
|
|
517
|
+
nodes_list_from_status = status_response.get("nodes")
|
|
518
|
+
tree_panel.build_metric_tree(nodes=nodes_list_from_status if nodes_list_from_status is not None else [])
|
|
502
519
|
tree_panel.set_unevaluated_node(node_id=eval_and_next_solution_response["solution_id"])
|
|
503
520
|
|
|
504
521
|
# Update the solution panels with the next solution and best solution (and score)
|
|
@@ -514,27 +531,25 @@ def main() -> None:
|
|
|
514
531
|
else:
|
|
515
532
|
best_solution_node = None
|
|
516
533
|
|
|
517
|
-
# Create a node for the current solution
|
|
518
534
|
current_solution_node = None
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
535
|
+
if status_response.get("nodes"):
|
|
536
|
+
for node_data in status_response["nodes"]:
|
|
537
|
+
if node_data["solution_id"] == eval_and_next_solution_response["solution_id"]:
|
|
538
|
+
current_solution_node = Node(
|
|
539
|
+
id=node_data["solution_id"],
|
|
540
|
+
parent_id=node_data["parent_id"],
|
|
541
|
+
code=node_data["code"],
|
|
542
|
+
metric=node_data["metric_value"],
|
|
543
|
+
is_buggy=node_data["is_buggy"],
|
|
544
|
+
)
|
|
528
545
|
if current_solution_node is None:
|
|
529
|
-
raise ValueError("Current solution node not found in
|
|
546
|
+
raise ValueError("Current solution node not found in nodes list from status response")
|
|
547
|
+
|
|
530
548
|
# Update the solution panels with the current and best solution
|
|
531
549
|
solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
|
|
532
550
|
current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
|
|
533
|
-
|
|
534
551
|
# Clear evaluation output since we are running a evaluation on a new solution
|
|
535
552
|
eval_output_panel.clear()
|
|
536
|
-
|
|
537
|
-
# Update displays with smooth transitions
|
|
538
553
|
smooth_update(
|
|
539
554
|
live=live,
|
|
540
555
|
layout=layout,
|
|
@@ -548,171 +563,152 @@ def main() -> None:
|
|
|
548
563
|
],
|
|
549
564
|
transition_delay=0.08, # Slightly longer delay for more noticeable transitions
|
|
550
565
|
)
|
|
551
|
-
|
|
552
|
-
# Run evaluation on the current solution
|
|
553
566
|
term_out = run_evaluation(eval_command=args.eval_command)
|
|
554
567
|
eval_output_panel.update(output=term_out)
|
|
555
|
-
|
|
556
|
-
# Update evaluation output with a smooth transition
|
|
557
568
|
smooth_update(
|
|
558
569
|
live=live,
|
|
559
570
|
layout=layout,
|
|
560
571
|
sections_to_update=[("eval_output", eval_output_panel.get_display())],
|
|
561
|
-
transition_delay=0.1,
|
|
572
|
+
transition_delay=0.1,
|
|
562
573
|
)
|
|
563
574
|
|
|
564
|
-
|
|
565
|
-
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
569
|
-
# Final evaluation report
|
|
570
|
-
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
|
|
571
|
-
session_id=session_id,
|
|
572
|
-
execution_output=term_out,
|
|
573
|
-
additional_instructions=current_additional_instructions,
|
|
574
|
-
api_keys=llm_api_keys,
|
|
575
|
-
timeout=timeout,
|
|
576
|
-
auth_headers=auth_headers,
|
|
577
|
-
)
|
|
578
|
-
|
|
579
|
-
# Update the progress bar
|
|
580
|
-
summary_panel.set_step(step=steps)
|
|
581
|
-
# Update the token counts
|
|
582
|
-
summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
|
|
583
|
-
# No need to update the plan panel since we have finished the optimization
|
|
584
|
-
# Get the optimization session status for
|
|
585
|
-
# the best solution, its score, and the history to plot the tree
|
|
586
|
-
status_response = get_optimization_session_status(
|
|
587
|
-
session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
|
|
588
|
-
)
|
|
589
|
-
# Build the metric tree
|
|
590
|
-
tree_panel.build_metric_tree(nodes=status_response["history"])
|
|
591
|
-
# No need to set any solution to unevaluated since we have finished the optimization
|
|
592
|
-
# and all solutions have been evaluated
|
|
593
|
-
# No neeed to update the current solution panel since we have finished the optimization
|
|
594
|
-
# We only need to update the best solution panel
|
|
595
|
-
# Figure out if we have a best solution so far
|
|
596
|
-
if status_response["best_result"] is not None:
|
|
597
|
-
best_solution_node = Node(
|
|
598
|
-
id=status_response["best_result"]["solution_id"],
|
|
599
|
-
parent_id=status_response["best_result"]["parent_id"],
|
|
600
|
-
code=status_response["best_result"]["code"],
|
|
601
|
-
metric=status_response["best_result"]["metric_value"],
|
|
602
|
-
is_buggy=status_response["best_result"]["is_buggy"],
|
|
575
|
+
if not user_stop_requested_flag:
|
|
576
|
+
# Re-read instructions from the original source (file path or string) BEFORE each suggest call
|
|
577
|
+
current_additional_instructions = read_additional_instructions(
|
|
578
|
+
additional_instructions=args.additional_instructions
|
|
603
579
|
)
|
|
604
|
-
|
|
605
|
-
|
|
606
|
-
|
|
607
|
-
|
|
608
|
-
|
|
609
|
-
|
|
610
|
-
|
|
611
|
-
|
|
612
|
-
if best_solution_node is not None and best_solution_node.metric is not None
|
|
613
|
-
else "[red] No valid solution found.[/]"
|
|
614
|
-
)
|
|
615
|
-
end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
|
|
616
|
-
end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
|
|
617
|
-
end_optimization_layout["best_solution"].update(best_solution_panel)
|
|
618
|
-
|
|
619
|
-
# Save optimization results
|
|
620
|
-
# If the best solution does not exist or is has not been measured at the end of the optimization
|
|
621
|
-
# save the original solution as the best solution
|
|
622
|
-
if best_solution_node is not None:
|
|
623
|
-
best_solution_code = best_solution_node.code
|
|
624
|
-
best_solution_score = best_solution_node.metric
|
|
625
|
-
else:
|
|
626
|
-
best_solution_code = None
|
|
627
|
-
best_solution_score = None
|
|
628
|
-
|
|
629
|
-
if best_solution_code is None or best_solution_score is None:
|
|
630
|
-
best_solution_content = f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_dir / f'step_0{source_fp.suffix}', is_json=False)}"
|
|
631
|
-
else:
|
|
632
|
-
# Format score for the comment
|
|
633
|
-
best_score_str = (
|
|
634
|
-
format_number(best_solution_score)
|
|
635
|
-
if best_solution_score is not None and isinstance(best_solution_score, (int, float))
|
|
636
|
-
else "N/A"
|
|
580
|
+
# Evaluate the final solution thats been generated
|
|
581
|
+
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
|
|
582
|
+
run_id=run_id,
|
|
583
|
+
execution_output=term_out,
|
|
584
|
+
additional_instructions=current_additional_instructions,
|
|
585
|
+
api_keys=llm_api_keys,
|
|
586
|
+
timeout=timeout,
|
|
587
|
+
auth_headers=auth_headers,
|
|
637
588
|
)
|
|
638
|
-
|
|
639
|
-
|
|
589
|
+
summary_panel.set_step(step=steps)
|
|
590
|
+
summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
|
|
591
|
+
status_response = get_optimization_run_status(
|
|
592
|
+
run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
|
|
640
593
|
)
|
|
594
|
+
# No need to update the plan panel since we have finished the optimization
|
|
595
|
+
# Get the optimization run status for
|
|
596
|
+
# the best solution, its score, and the history to plot the tree
|
|
597
|
+
nodes_list_from_status_final = status_response.get("nodes")
|
|
598
|
+
tree_panel.build_metric_tree(
|
|
599
|
+
nodes=nodes_list_from_status_final if nodes_list_from_status_final is not None else []
|
|
600
|
+
)
|
|
601
|
+
# No need to set any solution to unevaluated since we have finished the optimization
|
|
602
|
+
# and all solutions have been evaluated
|
|
603
|
+
# No neeed to update the current solution panel since we have finished the optimization
|
|
604
|
+
# We only need to update the best solution panel
|
|
605
|
+
# Figure out if we have a best solution so far
|
|
606
|
+
if status_response["best_result"] is not None:
|
|
607
|
+
best_solution_node = Node(
|
|
608
|
+
id=status_response["best_result"]["solution_id"],
|
|
609
|
+
parent_id=status_response["best_result"]["parent_id"],
|
|
610
|
+
code=status_response["best_result"]["code"],
|
|
611
|
+
metric=status_response["best_result"]["metric_value"],
|
|
612
|
+
is_buggy=status_response["best_result"]["is_buggy"],
|
|
613
|
+
)
|
|
614
|
+
else:
|
|
615
|
+
best_solution_node = None
|
|
616
|
+
solution_panels.update(current_node=None, best_node=best_solution_node)
|
|
617
|
+
_, best_solution_panel = solution_panels.get_display(current_step=steps)
|
|
618
|
+
# Update the end optimization layout
|
|
619
|
+
final_message = (
|
|
620
|
+
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']}[/] 🏆"
|
|
621
|
+
if best_solution_node is not None and best_solution_node.metric is not None
|
|
622
|
+
else "[red] No valid solution found.[/]"
|
|
623
|
+
)
|
|
624
|
+
end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
|
|
625
|
+
end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
|
|
626
|
+
end_optimization_layout["best_solution"].update(best_solution_panel)
|
|
627
|
+
|
|
628
|
+
# Save optimization results
|
|
629
|
+
# If the best solution does not exist or is has not been measured at the end of the optimization
|
|
630
|
+
# save the original solution as the best solution
|
|
631
|
+
if best_solution_node is not None:
|
|
632
|
+
best_solution_code = best_solution_node.code
|
|
633
|
+
best_solution_score = best_solution_node.metric
|
|
634
|
+
else:
|
|
635
|
+
best_solution_code = None
|
|
636
|
+
best_solution_score = None
|
|
641
637
|
|
|
642
|
-
|
|
643
|
-
|
|
644
|
-
|
|
645
|
-
|
|
646
|
-
|
|
647
|
-
|
|
648
|
-
|
|
649
|
-
|
|
650
|
-
|
|
651
|
-
|
|
638
|
+
if best_solution_code is None or best_solution_score is None:
|
|
639
|
+
best_solution_content = f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_dir / f'step_0{source_fp.suffix}', is_json=False)}"
|
|
640
|
+
else:
|
|
641
|
+
# Format score for the comment
|
|
642
|
+
best_score_str = (
|
|
643
|
+
format_number(best_solution_score)
|
|
644
|
+
if best_solution_score is not None and isinstance(best_solution_score, (int, float))
|
|
645
|
+
else "N/A"
|
|
646
|
+
)
|
|
647
|
+
best_solution_content = (
|
|
648
|
+
f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
|
|
649
|
+
)
|
|
650
|
+
# Save best solution to .runs/<run-id>/best.<extension>
|
|
651
|
+
write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
|
|
652
|
+
# write the best solution to the source file
|
|
653
|
+
write_to_path(fp=source_fp, content=best_solution_content)
|
|
654
|
+
# Mark as completed normally for the finally block
|
|
655
|
+
optimization_completed_normally = True
|
|
656
|
+
console.print(end_optimization_layout)
|
|
652
657
|
|
|
653
658
|
except Exception as e:
|
|
654
659
|
# Catch errors during the main optimization loop or setup
|
|
655
660
|
try:
|
|
656
|
-
error_message = e.response.json()["detail"]
|
|
661
|
+
error_message = e.response.json()["detail"]
|
|
657
662
|
except Exception:
|
|
658
|
-
error_message = str(e)
|
|
663
|
+
error_message = str(e)
|
|
659
664
|
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)
|
|
661
|
-
# console.print_exception(show_locals=False)
|
|
662
|
-
|
|
663
665
|
# Ensure optimization_completed_normally is False
|
|
664
666
|
optimization_completed_normally = False
|
|
665
|
-
|
|
666
|
-
# Prepare details for termination report
|
|
667
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
|
-
|
|
668
|
+
exit_code = 1
|
|
673
669
|
finally:
|
|
674
670
|
# This block runs whether the try block completed normally or raised an exception
|
|
675
|
-
|
|
676
671
|
# Stop heartbeat thread
|
|
677
672
|
stop_heartbeat_event.set()
|
|
678
673
|
if heartbeat_thread and heartbeat_thread.is_alive():
|
|
679
|
-
heartbeat_thread.join(timeout=2)
|
|
674
|
+
heartbeat_thread.join(timeout=2)
|
|
680
675
|
|
|
681
|
-
# Report final status if a
|
|
682
|
-
if
|
|
683
|
-
|
|
684
|
-
|
|
676
|
+
# Report final status if a run was started
|
|
677
|
+
if run_id:
|
|
678
|
+
final_status_update = "unknown"
|
|
679
|
+
final_reason_code = "unknown_termination"
|
|
685
680
|
final_details = None
|
|
686
|
-
|
|
687
681
|
if optimization_completed_normally:
|
|
688
|
-
|
|
689
|
-
|
|
682
|
+
final_status_update = "completed"
|
|
683
|
+
final_reason_code = "completed_successfully"
|
|
684
|
+
elif user_stop_requested_flag:
|
|
685
|
+
final_status_update = "terminated"
|
|
686
|
+
final_reason_code = "user_requested_stop"
|
|
687
|
+
final_details = "Run stopped by user request via dashboard."
|
|
690
688
|
else:
|
|
691
|
-
|
|
689
|
+
final_status_update = "error"
|
|
690
|
+
final_reason_code = "error_cli_internal"
|
|
692
691
|
if "error_details" in locals():
|
|
693
|
-
|
|
694
|
-
|
|
695
|
-
final_details =
|
|
696
|
-
|
|
697
|
-
|
|
698
|
-
|
|
692
|
+
final_details = locals()["error_details"]
|
|
693
|
+
elif "e" in locals() and isinstance(locals()["e"], Exception):
|
|
694
|
+
final_details = traceback.format_exc()
|
|
695
|
+
else:
|
|
696
|
+
final_details = "CLI terminated unexpectedly without a specific exception captured."
|
|
697
|
+
# Keep default 'unknown' if we somehow end up here without error/completion/signal
|
|
699
698
|
# Avoid reporting if terminated by signal handler (already reported)
|
|
700
699
|
# Check a flag or rely on status not being 'unknown'
|
|
701
|
-
if
|
|
700
|
+
if final_status_update != "unknown":
|
|
702
701
|
report_termination(
|
|
703
|
-
|
|
704
|
-
status_update=
|
|
705
|
-
reason=
|
|
702
|
+
run_id=run_id,
|
|
703
|
+
status_update=final_status_update,
|
|
704
|
+
reason=final_reason_code,
|
|
706
705
|
details=final_details,
|
|
707
|
-
auth_headers=
|
|
706
|
+
auth_headers=current_auth_headers_for_heartbeat,
|
|
708
707
|
)
|
|
709
|
-
|
|
710
|
-
|
|
711
|
-
|
|
712
|
-
|
|
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
|
|
708
|
+
if optimization_completed_normally:
|
|
709
|
+
sys.exit(0)
|
|
710
|
+
elif user_stop_requested_flag:
|
|
711
|
+
console.print("[yellow]Run terminated by user request.[/]")
|
|
718
712
|
sys.exit(0)
|
|
713
|
+
else:
|
|
714
|
+
sys.exit(locals().get("exit_code", 1))
|
weco/panels.py
CHANGED
|
@@ -11,9 +11,9 @@ from .__init__ import __dashboard_url__
|
|
|
11
11
|
|
|
12
12
|
|
|
13
13
|
class SummaryPanel:
|
|
14
|
-
"""Holds a summary of the optimization
|
|
14
|
+
"""Holds a summary of the optimization run."""
|
|
15
15
|
|
|
16
|
-
def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, runs_dir: str,
|
|
16
|
+
def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, runs_dir: str, run_id: str = None):
|
|
17
17
|
self.maximize = maximize
|
|
18
18
|
self.metric_name = metric_name
|
|
19
19
|
self.goal = ("Maximizing" if self.maximize else "Minimizing") + f" {self.metric_name}..."
|
|
@@ -22,7 +22,7 @@ class SummaryPanel:
|
|
|
22
22
|
self.total_steps = total_steps
|
|
23
23
|
self.model = model
|
|
24
24
|
self.runs_dir = runs_dir
|
|
25
|
-
self.
|
|
25
|
+
self.run_id = run_id if run_id is not None else "N/A"
|
|
26
26
|
self.dashboard_url = "N/A"
|
|
27
27
|
self.progress = Progress(
|
|
28
28
|
TextColumn("[progress.description]{task.description}"),
|
|
@@ -34,14 +34,14 @@ class SummaryPanel:
|
|
|
34
34
|
)
|
|
35
35
|
self.task_id = self.progress.add_task("", total=total_steps)
|
|
36
36
|
|
|
37
|
-
def
|
|
38
|
-
"""Set the
|
|
39
|
-
self.
|
|
40
|
-
self.set_dashboard_url(
|
|
37
|
+
def set_run_id(self, run_id: str):
|
|
38
|
+
"""Set the run ID."""
|
|
39
|
+
self.run_id = run_id
|
|
40
|
+
self.set_dashboard_url(run_id=run_id)
|
|
41
41
|
|
|
42
|
-
def set_dashboard_url(self,
|
|
42
|
+
def set_dashboard_url(self, run_id: str):
|
|
43
43
|
"""Set the dashboard URL."""
|
|
44
|
-
self.dashboard_url = f"{__dashboard_url__}/runs/{
|
|
44
|
+
self.dashboard_url = f"{__dashboard_url__}/runs/{run_id}"
|
|
45
45
|
|
|
46
46
|
def set_step(self, step: int):
|
|
47
47
|
"""Set the current step."""
|
|
@@ -70,7 +70,7 @@ class SummaryPanel:
|
|
|
70
70
|
summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
|
|
71
71
|
summary_table.add_row("")
|
|
72
72
|
# Log directory
|
|
73
|
-
summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.
|
|
73
|
+
summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.run_id}[/]")
|
|
74
74
|
summary_table.add_row("")
|
|
75
75
|
# Dashboard link
|
|
76
76
|
summary_table.add_row(f"[bold cyan]Dashboard:[/] [blue underline]{self.dashboard_url}[/]")
|
|
@@ -175,6 +175,9 @@ class MetricTreePanel:
|
|
|
175
175
|
|
|
176
176
|
def build_metric_tree(self, nodes: List[dict]):
|
|
177
177
|
"""Build the tree from the list of nodes."""
|
|
178
|
+
# Defensive: treat None as empty list
|
|
179
|
+
if nodes is None:
|
|
180
|
+
nodes = []
|
|
178
181
|
# First clear then tree
|
|
179
182
|
self.metric_tree.clear()
|
|
180
183
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: weco
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.20
|
|
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
|
|
@@ -29,6 +29,9 @@ Dynamic: license-file
|
|
|
29
29
|
[](https://docs.weco.ai/)
|
|
30
30
|
[](https://badge.fury.io/py/weco)
|
|
31
31
|
[](https://arxiv.org/abs/2502.13138)
|
|
32
|
+
[](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
|
|
33
|
+
|
|
34
|
+
`pip install weco`
|
|
32
35
|
|
|
33
36
|
</div>
|
|
34
37
|
|
|
@@ -98,9 +101,8 @@ pip install torch
|
|
|
98
101
|
weco run --source optimize.py \
|
|
99
102
|
--eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
|
|
100
103
|
--metric speedup \
|
|
101
|
-
--maximize
|
|
104
|
+
--goal maximize \
|
|
102
105
|
--steps 15 \
|
|
103
|
-
--model gemini-2.5-pro-exp-03-25 \
|
|
104
106
|
--additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."
|
|
105
107
|
```
|
|
106
108
|
|
|
@@ -108,18 +110,27 @@ weco run --source optimize.py \
|
|
|
108
110
|
|
|
109
111
|
---
|
|
110
112
|
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
|
116
|
-
|
|
|
117
|
-
|
|
|
118
|
-
|
|
|
119
|
-
|
|
|
120
|
-
|
|
|
121
|
-
|
|
122
|
-
|
|
113
|
+
### Arguments for `weco run`
|
|
114
|
+
|
|
115
|
+
**Required:**
|
|
116
|
+
|
|
117
|
+
| Argument | Description |
|
|
118
|
+
| :------------------ | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
|
119
|
+
| `-s, --source` | Path to the source code file that will be optimized (e.g., `optimize.py`). |
|
|
120
|
+
| `-c, --eval-command`| Command to run for evaluating the code in `--source`. This command should print the target `--metric` and its value to the terminal (stdout/stderr). See note below. |
|
|
121
|
+
| `-m, --metric` | The name of the metric you want to optimize (e.g., 'accuracy', 'speedup', 'loss'). This metric name should match what's printed by your `--eval-command`. |
|
|
122
|
+
| `-g, --goal` | `maximize`/`max` to maximize the `--metric` or `minimize`/`min` to minimize it. |
|
|
123
|
+
|
|
124
|
+
<br>
|
|
125
|
+
|
|
126
|
+
**Optional:**
|
|
127
|
+
|
|
128
|
+
| Argument | Description | Default |
|
|
129
|
+
| :----------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :------ |
|
|
130
|
+
| `-n, --steps` | Number of optimization steps (LLM iterations) to run. | 100 |
|
|
131
|
+
| `-M, --model` | Model identifier for the LLM to use (e.g., `gpt-4o`, `claude-3.5-sonnet`). | `o4-mini` when `OPENAI_API_KEY` is set; `claude-3-7-sonnet-20250219` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro-exp-03-25` when `GEMINI_API_KEY` is set (priority: `OPENAI_API_KEY` > `ANTHROPIC_API_KEY` > `GEMINI_API_KEY`). |
|
|
132
|
+
| `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM. | `None` |
|
|
133
|
+
| `-l, --log-dir` | Path to the directory to log intermediate steps and final optimization result. | `.runs/` |
|
|
123
134
|
|
|
124
135
|
---
|
|
125
136
|
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
weco/__init__.py,sha256=npWmRgLxfVK69GdyxIujnI87xqmPCBrZWxxAxL_QQOc,478
|
|
2
|
+
weco/api.py,sha256=xHCyQPto1Lv9QysiOFwVf5NnWDh6LBCNfPLyq-L7nys,5873
|
|
3
|
+
weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
|
|
4
|
+
weco/cli.py,sha256=e4h5bxeg2n95AlYXanfxLbcURWchjWTES2Kwx5AjKn0,36115
|
|
5
|
+
weco/panels.py,sha256=lsTHTh-XdYMH3ZV_WBteEcIt2hTWGGtqfUjGlYRHl70,13598
|
|
6
|
+
weco/utils.py,sha256=LVTBo3dduJmhlbotcYoUW2nLx6IRtKs4eDFR52Qltcg,5244
|
|
7
|
+
weco-0.2.20.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
|
|
8
|
+
weco-0.2.20.dist-info/METADATA,sha256=rK-Y9Q0zwaKUBS0bNZfUNvL82RXUiijVcIbW3i_IKKk,10955
|
|
9
|
+
weco-0.2.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
10
|
+
weco-0.2.20.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
|
|
11
|
+
weco-0.2.20.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
|
|
12
|
+
weco-0.2.20.dist-info/RECORD,,
|
weco-0.2.18.dist-info/RECORD
DELETED
|
@@ -1,12 +0,0 @@
|
|
|
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,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|