feat: add continuous optimization option(latest changes updated) (#76)

* fix: a little merge bug

* refactor: simplify optimization loop based on feedback

* fix: address review comments

* fix: remove redundant check for study.best_trials

* fix: restore comments

---------

Co-authored-by: Vinay Umrethe <vinayumrethe99@gmail.com>

src/heretic/main.py

@@ -467,6 +467,7 @@ def run():
         # defined in objective_wrapper above.
         pass
 
+    while True:
         # If no trials at all have been evaluated, the study must have been stopped
         # by pressing Ctrl+C while the first trial was running. In this case, we just
         # re-raise the interrupt to invoke the standard handler defined below.
@@ -504,6 +505,13 @@ def run():
             for trial in best_trials
         ]
 
+        choices.append(
+            Choice(
+                title="Continue optimization (run more trials)",
+                value="continue",
+            )
+        )
+
         choices.append(
             Choice(
                 title="None (exit program)",
@@ -527,8 +535,27 @@ def run():
             print()
             trial = prompt_select("Which trial do you want to use?", choices)
 
-        if trial is None or trial == "":
+            if trial == "continue":
+                while True:
+                    try:
+                        n_more_trials = int(
+                            prompt_text("How many more trials do you want to run?")
+                        )
+                        if n_more_trials > 0:
                             break
+                        print("[red]Please enter a number greater than 0.[/]")
+                    except ValueError:
+                        print("[red]Invalid input. Please enter a number.[/]")
+
+                settings.n_trials += n_more_trials
+                try:
+                    study.optimize(objective_wrapper, n_trials=n_more_trials)
+                except KeyboardInterrupt:
+                    pass
+                break
+
+            elif trial is None or trial == "":
+                return
 
             print()
             print(f"Restoring model from trial [bold]{trial.user_attrs['index']}[/]...")
@@ -556,7 +583,10 @@ def run():
                     ],
                 )
 
-            if action is None or action == "Nothing (return to trial selection menu)":
+                if (
+                    action is None
+                    or action == "Nothing (return to trial selection menu)"
+                ):
                     break
 
                 # All actions are wrapped in a try/except block so that if an error occurs,
@@ -696,7 +726,9 @@ def run():
 
                                     print("[bold]Assistant:[/] ", end="")
                                     response = model.stream_chat_response(chat)
-                                chat.append({"role": "assistant", "content": response})
+                                    chat.append(
+                                        {"role": "assistant", "content": response}
+                                    )
                                 except (KeyboardInterrupt, EOFError):
                                     # Ctrl+C/Ctrl+D
                                     break