Skip to content

Commit 8f9d7a4

Browse files
DhruvSrikanthDhruvSrikanth
authored
Merge pull request #13 from WecoAI/feature/request-cli-stop
Feature/request cli stop
2 parents 0eb5b9e + 141963f commit 8f9d7a4

File tree

1 file changed

+145
-106
lines changed

1 file changed

+145
-106
lines changed

‎weco/cli.py‎

Lines changed: 145 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,7 @@ def main() -> None:
314314

315315
session_id = None # Initialize session_id
316316
optimization_completed_normally = False # Flag for finally block
317+
user_stop_requested_flag = False # New flag for user-initiated stop
317318
# --- Check Authentication ---
318319
weco_api_key = load_weco_api_key()
319320
llm_api_keys = read_api_keys_from_env() # Read keys from client environment
@@ -504,6 +505,30 @@ def main() -> None:
504505
additional_instructions=args.additional_instructions
505506
)
506507

508+
# Get current session status BEFORE proceeding with suggestion/evaluation
509+
# This is where the CLI checks if it should stop
510+
if session_id: # Ensure session_id is available
511+
try:
512+
current_status_response = get_optimization_session_status(
513+
session_id=session_id,
514+
include_history=False, # Don't need full history here
515+
timeout=30, # Shorter timeout for status check
516+
auth_headers=auth_headers,
517+
)
518+
current_run_status = current_status_response.get("status")
519+
if current_run_status == "stopping":
520+
console.print("\n[bold yellow]Stop request received. Terminating run gracefully...[/]")
521+
user_stop_requested_flag = True
522+
break # Exit the optimization loop
523+
except requests.exceptions.RequestException as e:
524+
console.print(
525+
f"\n[bold red]Warning: Could not check session status: {e}. Continuing optimization...[/]"
526+
)
527+
except Exception as e: # Catch any other error during status check
528+
console.print(
529+
f"\n[bold red]Warning: Error checking session status: {e}. Continuing optimization...[/]"
530+
)
531+
507532
# Send feedback and get next suggestion
508533
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
509534
session_id=session_id,
@@ -599,94 +624,96 @@ def main() -> None:
599624
transition_delay=0.1, # Slightly longer delay for evaluation results
600625
)
601626

602-
# Re-read instructions from the original source (file path or string) BEFORE each suggest call
603-
current_additional_instructions = read_additional_instructions(
604-
additional_instructions=args.additional_instructions
605-
)
606-
607-
# Final evaluation report
608-
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
609-
session_id=session_id,
610-
execution_output=term_out,
611-
additional_instructions=current_additional_instructions,
612-
api_keys=llm_api_keys,
613-
timeout=timeout,
614-
auth_headers=auth_headers,
615-
)
627+
# If loop finished without user_stop_requested_flag
628+
if not user_stop_requested_flag:
629+
# Re-read instructions from the original source (file path or string) BEFORE each suggest call
630+
current_additional_instructions = read_additional_instructions(
631+
additional_instructions=args.additional_instructions
632+
)
616633

617-
# Update the progress bar
618-
summary_panel.set_step(step=steps)
619-
# Update the token counts
620-
summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
621-
# No need to update the plan panel since we have finished the optimization
622-
# Get the optimization session status for
623-
# the best solution, its score, and the history to plot the tree
624-
status_response = get_optimization_session_status(
625-
session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
626-
)
627-
# Build the metric tree
628-
tree_panel.build_metric_tree(nodes=status_response["history"])
629-
# No need to set any solution to unevaluated since we have finished the optimization
630-
# and all solutions have been evaluated
631-
# No neeed to update the current solution panel since we have finished the optimization
632-
# We only need to update the best solution panel
633-
# Figure out if we have a best solution so far
634-
if status_response["best_result"] is not None:
635-
best_solution_node = Node(
636-
id=status_response["best_result"]["solution_id"],
637-
parent_id=status_response["best_result"]["parent_id"],
638-
code=status_response["best_result"]["code"],
639-
metric=status_response["best_result"]["metric_value"],
640-
is_buggy=status_response["best_result"]["is_buggy"],
634+
# Final evaluation report
635+
eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
636+
session_id=session_id,
637+
execution_output=term_out,
638+
additional_instructions=current_additional_instructions,
639+
api_keys=llm_api_keys,
640+
timeout=timeout,
641+
auth_headers=auth_headers,
641642
)
642-
else:
643-
best_solution_node = None
644-
solution_panels.update(current_node=None, best_node=best_solution_node)
645-
_, best_solution_panel = solution_panels.get_display(current_step=steps)
646-
647-
# Update the end optimization layout
648-
final_message = (
649-
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']}[/] 🏆"
650-
if best_solution_node is not None and best_solution_node.metric is not None
651-
else "[red] No valid solution found.[/]"
652-
)
653-
end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
654-
end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
655-
end_optimization_layout["best_solution"].update(best_solution_panel)
656-
657-
# Save optimization results
658-
# If the best solution does not exist or is has not been measured at the end of the optimization
659-
# save the original solution as the best solution
660-
if best_solution_node is not None:
661-
best_solution_code = best_solution_node.code
662-
best_solution_score = best_solution_node.metric
663-
else:
664-
best_solution_code = None
665-
best_solution_score = None
666643

667-
if best_solution_code is None or best_solution_score is None:
668-
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)}"
669-
else:
670-
# Format score for the comment
671-
best_score_str = (
672-
format_number(best_solution_score)
673-
if best_solution_score is not None and isinstance(best_solution_score, (int, float))
674-
else "N/A"
644+
# Update the progress bar
645+
summary_panel.set_step(step=steps)
646+
# Update the token counts
647+
summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
648+
# No need to update the plan panel since we have finished the optimization
649+
# Get the optimization session status for
650+
# the best solution, its score, and the history to plot the tree
651+
status_response = get_optimization_session_status(
652+
session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
675653
)
676-
best_solution_content = (
677-
f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
654+
# Build the metric tree
655+
tree_panel.build_metric_tree(nodes=status_response["history"])
656+
# No need to set any solution to unevaluated since we have finished the optimization
657+
# and all solutions have been evaluated
658+
# No neeed to update the current solution panel since we have finished the optimization
659+
# We only need to update the best solution panel
660+
# Figure out if we have a best solution so far
661+
if status_response["best_result"] is not None:
662+
best_solution_node = Node(
663+
id=status_response["best_result"]["solution_id"],
664+
parent_id=status_response["best_result"]["parent_id"],
665+
code=status_response["best_result"]["code"],
666+
metric=status_response["best_result"]["metric_value"],
667+
is_buggy=status_response["best_result"]["is_buggy"],
668+
)
669+
else:
670+
best_solution_node = None
671+
solution_panels.update(current_node=None, best_node=best_solution_node)
672+
_, best_solution_panel = solution_panels.get_display(current_step=steps)
673+
674+
# Update the end optimization layout
675+
final_message = (
676+
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']}[/] 🏆"
677+
if best_solution_node is not None and best_solution_node.metric is not None
678+
else "[red] No valid solution found.[/]"
678679
)
680+
end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
681+
end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
682+
end_optimization_layout["best_solution"].update(best_solution_panel)
683+
684+
# Save optimization results
685+
# If the best solution does not exist or is has not been measured at the end of the optimization
686+
# save the original solution as the best solution
687+
if best_solution_node is not None:
688+
best_solution_code = best_solution_node.code
689+
best_solution_score = best_solution_node.metric
690+
else:
691+
best_solution_code = None
692+
best_solution_score = None
679693

680-
# Save best solution to .runs/<session-id>/best.<extension>
681-
write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
694+
if best_solution_code is None or best_solution_score is None:
695+
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)}"
696+
else:
697+
# Format score for the comment
698+
best_score_str = (
699+
format_number(best_solution_score)
700+
if best_solution_score is not None and isinstance(best_solution_score, (int, float))
701+
else "N/A"
702+
)
703+
best_solution_content = (
704+
f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
705+
)
706+
707+
# Save best solution to .runs/<session-id>/best.<extension>
708+
write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
682709

683-
# write the best solution to the source file
684-
write_to_path(fp=source_fp, content=best_solution_content)
710+
# write the best solution to the source file
711+
write_to_path(fp=source_fp, content=best_solution_content)
685712

686-
# Mark as completed normally for the finally block
687-
optimization_completed_normally = True
713+
# Mark as completed normally for the finally block
714+
optimization_completed_normally = True # Only set if loop completes all steps
688715

689-
console.print(end_optimization_layout)
716+
console.print(end_optimization_layout) # Moved inside the if
690717

691718
except Exception as e:
692719
# Catch errors during the main optimization loop or setup
@@ -718,39 +745,51 @@ def main() -> None:
718745

719746
# Report final status if a session was started
720747
if session_id:
721-
final_status = "unknown"
722-
final_reason = "unknown_termination"
748+
final_status_update = "unknown"
749+
final_reason_code = "unknown_termination"
723750
final_details = None
724751

725-
if optimization_completed_normally:
726-
final_status = "completed"
727-
final_reason = "completed_successfully"
728-
else:
729-
# If an exception was caught and we have details
752+
if optimization_completed_normally: # All steps completed
753+
final_status_update = "completed"
754+
final_reason_code = "completed_successfully"
755+
elif user_stop_requested_flag: # Stopped by user request
756+
final_status_update = "terminated"
757+
final_reason_code = "user_requested_stop"
758+
final_details = "Run stopped by user request via dashboard."
759+
else: # Any other non-normal termination (e.g., CLI error)
760+
final_status_update = "error"
761+
final_reason_code = "error_cli_internal"
762+
# Use error_details from the existing except block if available
730763
if "error_details" in locals():
731-
final_status = "error"
732-
final_reason = "error_cli_internal"
733-
final_details = error_details
734-
# else: # Should have been handled by signal handler if terminated by user
735-
# Keep default 'unknown' if we somehow end up here without error/completion/signal
736-
737-
# Avoid reporting if terminated by signal handler (already reported)
738-
# Check a flag or rely on status not being 'unknown'
739-
if final_status != "unknown":
764+
final_details = locals()["error_details"]
765+
elif "e" in locals() and isinstance(locals()["e"], Exception): # Fallback if e is somehow there
766+
final_details = traceback.format_exc()
767+
else:
768+
final_details = "CLI terminated unexpectedly without a specific exception captured."
769+
770+
# The signal_handler will call report_termination with its own reasons.
771+
# This `finally` block's report_termination is for loop completion,
772+
# user_requested_stop, or internal CLI errors not caught by signals.
773+
# We assume signal_handler calls sys.exit, so this block might not fully
774+
# execute if a signal terminates the process abruptly *before* this finally.
775+
# However, for user_requested_stop detected by the loop, this will run.
776+
777+
if final_status_update != "unknown":
740778
report_termination(
741779
session_id=session_id,
742-
status_update=final_status,
743-
reason=final_reason,
780+
status_update=final_status_update,
781+
reason=final_reason_code,
744782
details=final_details,
745-
auth_headers=auth_headers,
783+
auth_headers=current_auth_headers_for_heartbeat,
746784
)
747785

748-
# Ensure proper exit code if an error occurred
749-
if not optimization_completed_normally and "exit_code" in locals() and exit_code != 0:
750-
sys.exit(exit_code)
751-
elif not optimization_completed_normally:
752-
# Generic error exit if no specific code was set but try block failed
753-
sys.exit(1)
754-
else:
755-
# Normal exit
786+
# Exit code logic
787+
if optimization_completed_normally:
756788
sys.exit(0)
789+
elif user_stop_requested_flag:
790+
console.print("[yellow]Run terminated by user request.[/]")
791+
sys.exit(0) # Graceful exit for user stop
792+
else:
793+
# If an error occurred and was caught by the `except Exception as e` block
794+
# exit_code might have been set by the existing except block.
795+
sys.exit(locals().get("exit_code", 1))

0 commit comments

Comments
 (0)