weco 0.3.4__py3-none-any.whl → 0.3.6__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/optimizer.py CHANGED
@@ -127,14 +127,15 @@ def execute_optimization(
127
127
  eval_command: str,
128
128
  metric: str,
129
129
  goal: str, # "maximize" or "minimize"
130
+ model: str,
130
131
  steps: int = 100,
131
- model: Optional[str] = None,
132
132
  log_dir: str = ".runs",
133
133
  additional_instructions: Optional[str] = None,
134
134
  console: Optional[Console] = None,
135
135
  eval_timeout: Optional[int] = None,
136
136
  save_logs: bool = False,
137
137
  apply_change: bool = False,
138
+ api_keys: Optional[dict[str, str]] = None,
138
139
  ) -> bool:
139
140
  """
140
141
  Execute the core optimization logic.
@@ -144,17 +145,25 @@ def execute_optimization(
144
145
  """
145
146
  if console is None:
146
147
  console = Console()
147
-
148
148
  # Global variables for this optimization run
149
149
  heartbeat_thread = None
150
150
  stop_heartbeat_event = threading.Event()
151
151
  current_run_id_for_heartbeat = None
152
152
  current_auth_headers_for_heartbeat = {}
153
+ live_ref = None # Reference to the Live object for the optimization run
154
+
155
+ best_solution_code = None
156
+ original_source_code = None
153
157
 
154
158
  # --- Signal Handler for this optimization run ---
155
159
  def signal_handler(signum, frame):
160
+ nonlocal live_ref
161
+
162
+ if live_ref is not None:
163
+ live_ref.stop() # Stop the live update loop so that messages are printed to the console
164
+
156
165
  signal_name = signal.Signals(signum).name
157
- console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]")
166
+ console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]\n")
158
167
 
159
168
  # Stop heartbeat thread
160
169
  stop_heartbeat_event.set()
@@ -170,7 +179,7 @@ def execute_optimization(
170
179
  details=f"Process terminated by signal {signal_name} ({signum}).",
171
180
  auth_headers=current_auth_headers_for_heartbeat,
172
181
  )
173
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]")
182
+ console.print(f"[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]\n")
174
183
 
175
184
  # Exit gracefully
176
185
  sys.exit(0)
@@ -183,8 +192,6 @@ def execute_optimization(
183
192
  optimization_completed_normally = False
184
193
  user_stop_requested_flag = False
185
194
 
186
- best_solution_code = None
187
- original_source_code = None # Make available to the finally block
188
195
  try:
189
196
  # --- Login/Authentication Handling (now mandatory) ---
190
197
  weco_api_key, auth_headers = handle_authentication(console)
@@ -197,11 +204,6 @@ def execute_optimization(
197
204
  # --- Process Parameters ---
198
205
  maximize = goal.lower() in ["maximize", "max"]
199
206
 
200
- # Determine the model to use
201
- if model is None:
202
- # Default to o4-mini with credit-based billing
203
- model = "o4-mini"
204
-
205
207
  code_generator_config = {"model": model}
206
208
  evaluator_config = {"model": model, "include_analysis": True}
207
209
  search_policy_config = {
@@ -239,6 +241,7 @@ def execute_optimization(
239
241
  save_logs=save_logs,
240
242
  log_dir=log_dir,
241
243
  auth_headers=auth_headers,
244
+ api_keys=api_keys,
242
245
  )
243
246
  # Indicate the endpoint failed to return a response and the optimization was unsuccessful
244
247
  if run_response is None:
@@ -256,6 +259,7 @@ def execute_optimization(
256
259
  # --- Live Update Loop ---
257
260
  refresh_rate = 4
258
261
  with Live(layout, refresh_per_second=refresh_rate) as live:
262
+ live_ref = live
259
263
  # Define the runs directory (.runs/<run-id>) to store logs and results
260
264
  runs_dir = pathlib.Path(log_dir) / run_id
261
265
  runs_dir.mkdir(parents=True, exist_ok=True)
@@ -364,7 +368,12 @@ def execute_optimization(
364
368
 
365
369
  # Send feedback and get next suggestion
366
370
  eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
367
- console=console, step=step, run_id=run_id, execution_output=term_out, auth_headers=auth_headers
371
+ console=console,
372
+ step=step,
373
+ run_id=run_id,
374
+ execution_output=term_out,
375
+ auth_headers=auth_headers,
376
+ api_keys=api_keys,
368
377
  )
369
378
  # Save next solution (.runs/<run-id>/step_<step>.<extension>)
370
379
  write_to_path(fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"])
@@ -387,6 +396,16 @@ def execute_optimization(
387
396
  status_response=status_response, solution_id=eval_and_next_solution_response["solution_id"]
388
397
  )
389
398
 
399
+ # Set best solution and save optimization results
400
+ try:
401
+ best_solution_code = best_solution_node.code
402
+ except AttributeError:
403
+ # Can happen if the code was buggy
404
+ best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
405
+
406
+ # Save best solution to .runs/<run-id>/best.<extension>
407
+ write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
408
+
390
409
  # Update the solution panels with the current and best solution
391
410
  solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
392
411
  current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
@@ -428,7 +447,12 @@ def execute_optimization(
428
447
  if not user_stop_requested_flag:
429
448
  # Evaluate the final solution thats been generated
430
449
  eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
431
- console=console, step=steps, run_id=run_id, execution_output=term_out, auth_headers=auth_headers
450
+ console=console,
451
+ step=steps,
452
+ run_id=run_id,
453
+ execution_output=term_out,
454
+ auth_headers=auth_headers,
455
+ api_keys=api_keys,
432
456
  )
433
457
  summary_panel.set_step(step=steps)
434
458
  status_response = get_optimization_run_status(
@@ -443,10 +467,13 @@ def execute_optimization(
443
467
  )
444
468
  # No need to set any solution to unevaluated since we have finished the optimization
445
469
  # and all solutions have been evaluated
446
- # No neeed to update the current solution panel since we have finished the optimization
470
+ # No need to update the current solution panel since we have finished the optimization
447
471
  # We only need to update the best solution panel
448
472
  # Figure out if we have a best solution so far
449
473
  best_solution_node = get_best_node_from_status(status_response=status_response)
474
+ best_solution_code = best_solution_node.code
475
+ # Save best solution to .runs/<run-id>/best.<extension>
476
+ write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
450
477
  solution_panels.update(current_node=None, best_node=best_solution_node)
451
478
  _, best_solution_panel = solution_panels.get_display(current_step=steps)
452
479
  # Update the end optimization layout
@@ -459,17 +486,6 @@ def execute_optimization(
459
486
  end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
460
487
  end_optimization_layout["best_solution"].update(best_solution_panel)
461
488
 
462
- # Save optimization results
463
- # If the best solution does not exist or is has not been measured at the end of the optimization
464
- # save the original solution as the best solution
465
- try:
466
- best_solution_code = best_solution_node.code
467
- except AttributeError:
468
- best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
469
-
470
- # Save best solution to .runs/<run-id>/best.<extension>
471
- write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
472
-
473
489
  # Mark as completed normally for the finally block
474
490
  optimization_completed_normally = True
475
491
  live.update(end_optimization_layout)
@@ -481,7 +497,7 @@ def execute_optimization(
481
497
  except Exception:
482
498
  error_message = str(e)
483
499
  console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Optimization Error", border_style="red"))
484
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]")
500
+ console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]\n")
485
501
  # Ensure optimization_completed_normally is False
486
502
  optimization_completed_normally = False
487
503
  finally:
@@ -508,20 +524,16 @@ def execute_optimization(
508
524
  else "CLI terminated unexpectedly without a specific exception captured."
509
525
  )
510
526
 
511
- # raise Exception(best_solution_code, original_source_code)
512
527
  if best_solution_code and best_solution_code != original_source_code:
513
528
  # Determine whether to apply: automatically if --apply-change is set, otherwise ask user
514
- should_apply = apply_change or summary_panel.ask_user_feedback(
515
- live=live,
516
- layout=end_optimization_layout,
517
- question="Would you like to apply the best solution to the source file?",
518
- default=True,
529
+ should_apply = apply_change or Confirm.ask(
530
+ "Would you like to apply the best solution to the source file?", default=True
519
531
  )
520
532
  if should_apply:
521
533
  write_to_path(fp=source_fp, content=best_solution_code)
522
- console.print("[green]Best solution applied to the source file.[/]\n")
534
+ console.print("\n[green]Best solution applied to the source file.[/]\n")
523
535
  else:
524
- console.print("[green]A better solution was not found. No changes to apply.[/]\n")
536
+ console.print("\n[green]A better solution was not found. No changes to apply.[/]\n")
525
537
 
526
538
  report_termination(
527
539
  run_id=run_id,
@@ -534,12 +546,14 @@ def execute_optimization(
534
546
  # Handle exit
535
547
  if user_stop_requested_flag:
536
548
  console.print("[yellow]Run terminated by user request.[/]")
537
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]")
549
+ console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]\n")
538
550
 
539
551
  return optimization_completed_normally or user_stop_requested_flag
540
552
 
541
553
 
542
- def resume_optimization(run_id: str, console: Optional[Console] = None, apply_change: bool = False) -> bool:
554
+ def resume_optimization(
555
+ run_id: str, console: Optional[Console] = None, apply_change: bool = False, api_keys: Optional[dict[str, str]] = None
556
+ ) -> bool:
543
557
  """Resume an interrupted run from the most recent node and continue optimization."""
544
558
  if console is None:
545
559
  console = Console()
@@ -549,11 +563,19 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
549
563
  stop_heartbeat_event = threading.Event()
550
564
  current_run_id_for_heartbeat = None
551
565
  current_auth_headers_for_heartbeat = {}
566
+ live_ref = None # Reference to the Live object for the optimization run
567
+
568
+ best_solution_code = None
569
+ original_source_code = None
552
570
 
553
571
  # Signal handler for this optimization run
554
572
  def signal_handler(signum, frame):
573
+ nonlocal live_ref
574
+ if live_ref is not None:
575
+ live_ref.stop() # Stop the live update loop so that messages are printed to the console
576
+
555
577
  signal_name = signal.Signals(signum).name
556
- console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]")
578
+ console.print(f"\n[bold yellow]Termination signal ({signal_name}) received. Shutting down...[/]\n")
557
579
  stop_heartbeat_event.set()
558
580
  if heartbeat_thread and heartbeat_thread.is_alive():
559
581
  heartbeat_thread.join(timeout=2)
@@ -565,7 +587,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
565
587
  details=f"Process terminated by signal {signal_name} ({signum}).",
566
588
  auth_headers=current_auth_headers_for_heartbeat,
567
589
  )
568
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]")
590
+ console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]\n")
569
591
  sys.exit(0)
570
592
 
571
593
  # Set up signal handlers for this run
@@ -575,9 +597,6 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
575
597
  optimization_completed_normally = False
576
598
  user_stop_requested_flag = False
577
599
 
578
- best_solution_code = None
579
- original_source_code = None
580
-
581
600
  try:
582
601
  # --- Login/Authentication Handling (now mandatory) ---
583
602
  weco_api_key, auth_headers = handle_authentication(console)
@@ -659,9 +678,8 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
659
678
  source_fp.parent.mkdir(parents=True, exist_ok=True)
660
679
  # Store the original content to restore after each evaluation
661
680
  original_source_code = read_from_path(fp=source_fp, is_json=False) if source_fp.exists() else ""
662
-
681
+ # The code to restore is the code from the last step of the previous run
663
682
  code_to_restore = resume_resp.get("code") or resume_resp.get("source_code") or ""
664
- write_to_path(fp=source_fp, content=code_to_restore)
665
683
 
666
684
  # Prepare UI panels
667
685
  summary_panel = SummaryPanel(
@@ -687,12 +705,25 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
687
705
  best_solution_node = get_best_node_from_status(status_response=status)
688
706
  current_solution_node = get_node_from_status(status_response=status, solution_id=resume_resp.get("solution_id"))
689
707
 
708
+ # If there's no best solution yet (baseline evaluation didn't complete),
709
+ # mark the current node as unevaluated so the tree renders correctly
710
+ if best_solution_node is None:
711
+ tree_panel.set_unevaluated_node(node_id=resume_resp.get("solution_id"))
712
+
690
713
  # Ensure runs dir exists
691
714
  runs_dir = pathlib.Path(log_dir) / resume_resp["run_id"]
692
715
  runs_dir.mkdir(parents=True, exist_ok=True)
693
716
  # Persist last step's code into logs as step_<current_step>
694
717
  write_to_path(fp=runs_dir / f"step_{current_step}{source_fp.suffix}", content=code_to_restore)
695
718
 
719
+ # Initialize best solution code
720
+ try:
721
+ best_solution_code = best_solution_node.code
722
+ except AttributeError:
723
+ # Edge case: best solution node is not available.
724
+ # This can happen if the user has cancelled the run before even running the baseline solution
725
+ pass # Leave best solution code as None
726
+
696
727
  # Start Heartbeat Thread
697
728
  stop_heartbeat_event.clear()
698
729
  heartbeat_thread = HeartbeatSender(resume_resp["run_id"], auth_headers, stop_heartbeat_event)
@@ -705,6 +736,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
705
736
  # --- Live UI ---
706
737
  refresh_rate = 4
707
738
  with Live(layout, refresh_per_second=refresh_rate) as live:
739
+ live_ref = live
708
740
  # Initial panels
709
741
  current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=current_step)
710
742
  # Use backend-provided execution output only (no fallback)
@@ -769,6 +801,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
769
801
  run_id=resume_resp["run_id"],
770
802
  execution_output=term_out,
771
803
  auth_headers=auth_headers,
804
+ api_keys=api_keys,
772
805
  )
773
806
 
774
807
  # Save next solution to logs
@@ -788,6 +821,16 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
788
821
  status_response=status_response, solution_id=eval_and_next_solution_response["solution_id"]
789
822
  )
790
823
 
824
+ # Set best solution and save optimization results
825
+ try:
826
+ best_solution_code = best_solution_node.code
827
+ except AttributeError:
828
+ # Can happen if the code was buggy
829
+ best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
830
+
831
+ # Save best solution to .runs/<run-id>/best.<extension>
832
+ write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
833
+
791
834
  solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
792
835
  current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
793
836
  eval_output_panel.clear()
@@ -830,6 +873,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
830
873
  run_id=resume_resp["run_id"],
831
874
  execution_output=term_out,
832
875
  auth_headers=auth_headers,
876
+ api_keys=api_keys,
833
877
  )
834
878
  summary_panel.set_step(step=total_steps)
835
879
  status_response = get_optimization_run_status(
@@ -839,6 +883,10 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
839
883
  tree_panel.build_metric_tree(nodes=nodes_final)
840
884
  # Best solution panel and final message
841
885
  best_solution_node = get_best_node_from_status(status_response=status_response)
886
+ best_solution_code = best_solution_node.code
887
+ # Save best solution to .runs/<run-id>/best.<extension>
888
+ write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
889
+
842
890
  solution_panels.update(current_node=None, best_node=best_solution_node)
843
891
  _, best_solution_panel = solution_panels.get_display(current_step=total_steps)
844
892
  final_message = (
@@ -850,14 +898,6 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
850
898
  end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
851
899
  end_optimization_layout["best_solution"].update(best_solution_panel)
852
900
 
853
- # Save best
854
- try:
855
- best_solution_code = best_solution_node.code
856
- except AttributeError:
857
- best_solution_code = read_from_path(fp=runs_dir / f"step_0{source_fp.suffix}", is_json=False)
858
-
859
- write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_code)
860
-
861
901
  optimization_completed_normally = True
862
902
  live.update(end_optimization_layout)
863
903
 
@@ -867,7 +907,7 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
867
907
  except Exception:
868
908
  error_message = str(e)
869
909
  console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Optimization Error", border_style="red"))
870
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]")
910
+ console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]\n")
871
911
  optimization_completed_normally = False
872
912
  finally:
873
913
  signal.signal(signal.SIGINT, original_sigint_handler)
@@ -896,17 +936,14 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
896
936
  )
897
937
 
898
938
  if best_solution_code and best_solution_code != original_source_code:
899
- should_apply = apply_change or summary_panel.ask_user_feedback(
900
- live=live,
901
- layout=end_optimization_layout,
902
- question="Would you like to apply the best solution to the source file?",
903
- default=True,
939
+ should_apply = apply_change or Confirm.ask(
940
+ "Would you like to apply the best solution to the source file?", default=True
904
941
  )
905
942
  if should_apply:
906
943
  write_to_path(fp=source_fp, content=best_solution_code)
907
- console.print("[green]Best solution applied to the source file.[/]\n")
944
+ console.print("\n[green]Best solution applied to the source file.[/]\n")
908
945
  else:
909
- console.print("[green]A better solution was not found. No changes to apply.[/]\n")
946
+ console.print("\n[green]A better solution was not found. No changes to apply.[/]\n")
910
947
 
911
948
  report_termination(
912
949
  run_id=run_id,
@@ -917,5 +954,5 @@ def resume_optimization(run_id: str, console: Optional[Console] = None, apply_ch
917
954
  )
918
955
  if user_stop_requested_flag:
919
956
  console.print("[yellow]Run terminated by user request.[/]")
920
- console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]")
957
+ console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {run_id}[/]\n")
921
958
  return optimization_completed_normally or user_stop_requested_flag
weco/panels.py CHANGED
@@ -6,9 +6,7 @@ from rich.panel import Panel
6
6
  from rich.syntax import Syntax
7
7
  from rich import box
8
8
  from rich.console import Console
9
- from rich.live import Live
10
9
 
11
- from rich.prompt import Confirm
12
10
  from typing import Dict, List, Optional, Union, Tuple
13
11
  from pathlib import Path
14
12
  from .__init__ import __dashboard_url__
@@ -74,25 +72,6 @@ class SummaryPanel:
74
72
  """Clear the thinking content."""
75
73
  self.thinking_content = ""
76
74
 
77
- def ask_user_feedback(self, live: Live, layout: Layout, question: str, default: bool = True) -> bool:
78
- """
79
- Ask a yes/no question while keeping the main layout fixed.
80
- Uses Rich's Confirm for a clean user experience.
81
- """
82
- # Stop live updates temporarily to prevent layout from moving
83
- live.stop()
84
-
85
- try:
86
- # Use Rich's built-in Confirm
87
- result = Confirm.ask(question, default=default)
88
- except (KeyboardInterrupt, EOFError):
89
- result = default
90
- finally:
91
- # Resume live updates
92
- live.start()
93
-
94
- return result
95
-
96
75
  def get_display(self, final_message: Optional[str] = None) -> Panel:
97
76
  """Return a Rich panel summarising the current run."""
98
77
  # ───────────────────── summary grid ──────────────────────
weco/utils.py CHANGED
@@ -9,13 +9,26 @@ from rich.panel import Panel
9
9
  import pathlib
10
10
  import requests
11
11
  from packaging.version import parse as parse_version
12
- from .constants import TRUNCATION_THRESHOLD, TRUNCATION_KEEP_LENGTH, DEFAULT_MODEL, SUPPORTED_FILE_EXTENSIONS
12
+ from .constants import TRUNCATION_THRESHOLD, TRUNCATION_KEEP_LENGTH, SUPPORTED_FILE_EXTENSIONS, DEFAULT_MODELS
13
13
 
14
14
 
15
- # Env/arg helper functions
16
- def determine_model_for_onboarding() -> str:
17
- """Determine which model to use for onboarding chatbot. Defaults to o4-mini."""
18
- return DEFAULT_MODEL
15
+ class UnrecognizedAPIKeysError(Exception):
16
+ """Exception raised when unrecognized API keys are provided."""
17
+
18
+ def __init__(self, api_keys: dict[str, str]):
19
+ self.api_keys = api_keys
20
+ providers = {provider for provider, _ in DEFAULT_MODELS}
21
+ super().__init__(
22
+ f"Unrecognized API key provider in {set(api_keys.keys())}. Supported providers: {', '.join(providers)}"
23
+ )
24
+
25
+
26
+ class DefaultModelNotFoundError(Exception):
27
+ """Exception raised when no default model is found for the API keys."""
28
+
29
+ def __init__(self, api_keys: dict[str, str]):
30
+ self.api_keys = api_keys
31
+ super().__init__(f"No default model found for any of the provided API keys: {set(api_keys.keys())}")
19
32
 
20
33
 
21
34
  def read_additional_instructions(additional_instructions: str | None) -> str | None:
@@ -218,3 +231,18 @@ def check_for_cli_updates():
218
231
  except Exception:
219
232
  # Catch any other unexpected error during the check
220
233
  pass
234
+
235
+
236
+ def get_default_model(api_keys: dict[str, str] | None = None) -> str:
237
+ """Determine the default model to use based on the API keys."""
238
+ providers = {provider for provider, _ in DEFAULT_MODELS}
239
+ if api_keys and not all(provider in providers for provider in api_keys.keys()):
240
+ raise UnrecognizedAPIKeysError(api_keys)
241
+
242
+ if api_keys:
243
+ for provider, model in DEFAULT_MODELS:
244
+ if provider in api_keys:
245
+ return model
246
+ # Should never happen, but just in case
247
+ raise DefaultModelNotFoundError(api_keys)
248
+ return DEFAULT_MODELS[0][1]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: weco
3
- Version: 0.3.4
3
+ Version: 0.3.6
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:
@@ -224,6 +224,7 @@ Provides-Extra: dev
224
224
  Requires-Dist: ruff; extra == "dev"
225
225
  Requires-Dist: build; extra == "dev"
226
226
  Requires-Dist: setuptools_scm; extra == "dev"
227
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
227
228
  Dynamic: license-file
228
229
 
229
230
  <div align="center">
@@ -323,6 +324,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
323
324
  | `--eval-timeout` | Timeout in seconds for each step in evaluation. | No timeout (unlimited) | `--eval-timeout 3600` |
324
325
  | `--save-logs` | Save execution output from each optimization step to disk. Creates timestamped directories with raw output files and a JSONL index for tracking execution history. | `False` | `--save-logs` |
325
326
  | `--apply-change` | Automatically apply the best solution to the source file without prompting. | `False` | `--apply-change` |
327
+ | `--api-key` | API keys for LLM providers (BYOK). Format: `provider=key`. Can specify multiple providers. | `None` | `--api-key openai=sk-xxx` |
326
328
 
327
329
  ---
328
330
 
@@ -377,6 +379,7 @@ Arguments for `weco resume`:
377
379
  |----------|-------------|---------|
378
380
  | `run-id` | The UUID of the run to resume (shown at the start of each run) | `0002e071-1b67-411f-a514-36947f0c4b31` |
379
381
  | `--apply-change` | Automatically apply the best solution to the source file without prompting | `--apply-change` |
382
+ | `--api-key` | (Optional) API keys for LLM providers (BYOK). Format: `provider=key` | `--api-key openai=sk-xxx` |
380
383
 
381
384
  Notes:
382
385
  - Works only for interrupted runs (status: `error`, `terminated`, etc.).
@@ -0,0 +1,15 @@
1
+ weco/__init__.py,sha256=ClO0uT6GKOA0iSptvP0xbtdycf0VpoPTq37jHtvlhtw,303
2
+ weco/api.py,sha256=xVVRk1pj9jpjTphaInkkAhjqhgFP2-6zHT_V-5Du1Fc,13629
3
+ weco/auth.py,sha256=O31Hoj-Loi8DWJJG2LfeWgUMuNqAUeGDpd2ZGjA9Ah0,9997
4
+ weco/cli.py,sha256=Mtkv3rE1rQLdeoVydn30EUi1ki3Cyu45Q3cONnQH4QY,11210
5
+ weco/constants.py,sha256=Wt1VI76smsU0Gsg-HSB9e3V1ZjjsMYXJJg2G9Rb-nQc,519
6
+ weco/credits.py,sha256=C08x-TRcLg3ccfKqMGNRY7zBn7t3r7LZ119bxgfztaI,7629
7
+ weco/optimizer.py,sha256=2qYweESOAer26gjjhu4dg01XmuhtA2nMrJuijaxizzE,45492
8
+ weco/panels.py,sha256=POHt0MdRKDykwUJYXcry92O41lpB9gxna55wFI9abWU,16272
9
+ weco/utils.py,sha256=v_rvgw-ktRoXrpPA2copngI8QDCB8UXmbiN-wAiYvEE,9450
10
+ weco-0.3.6.dist-info/licenses/LICENSE,sha256=9LUfoGHjLPtak2zps2kL2tm65HAZIICx_FbLaRuS4KU,11337
11
+ weco-0.3.6.dist-info/METADATA,sha256=-BKVxw7j_Qe5ksblcmoVF_NVeDDIcWTg-CsUUaadAgs,29835
12
+ weco-0.3.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
13
+ weco-0.3.6.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
14
+ weco-0.3.6.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
15
+ weco-0.3.6.dist-info/RECORD,,