From 1c82514a8c4549f27b3c84974850d05e2f25a6cc Mon Sep 17 00:00:00 2001
From: Yiannis Charalambous <yiannis128@hotmail.com>
Date: Wed, 3 Apr 2024 18:57:52 +0100
Subject: [PATCH] FixCodeCommand: Updated notice messages and bug fix

Fixed bug where if verification is successful then the get_esbmc_output_formatted would return a stack trace.
---
 esbmc_ai/commands/fix_code_command.py | 41 +++++++++++++++------------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/esbmc_ai/commands/fix_code_command.py b/esbmc_ai/commands/fix_code_command.py
index fb103a4..9ed2edc 100644
--- a/esbmc_ai/commands/fix_code_command.py
+++ b/esbmc_ai/commands/fix_code_command.py
@@ -1,6 +1,5 @@
 # Author: Yiannis Charalambous
 
-from os import get_terminal_size
 import sys
 from typing import Any, Tuple
 from typing_extensions import override
@@ -40,12 +39,13 @@ def __init__(self) -> None:
     @override
     def execute(self, **kwargs: Any) -> Tuple[bool, str]:
         def print_raw_conversation() -> None:
-            print("Notice: Printing raw conversation...")
+            print_horizontal_line(0)
+            print("ESBMC-AI Notice: Printing raw conversation...")
             all_messages = solution_generator._system_messages.copy()
             all_messages.extend(solution_generator.messages.copy())
             messages: list[str] = [f"{msg.type}: {msg.content}" for msg in all_messages]
             print("\n" + "\n\n".join(messages))
-            print("Notice: End of conversation")
+            print("ESBMC-AI Notice: End of raw conversation")
 
         file_name: str = kwargs["file_name"]
         source_code: str = kwargs["source_code"]
@@ -119,34 +119,38 @@ def print_raw_conversation() -> None:
             )
             self.anim.stop()
 
-            # TODO Move this process into Solution Generator since have (beginning) is done
-            # inside, and the other half is done here.
-            try:
-                esbmc_output = get_esbmc_output_formatted(
-                    esbmc_output_type=config.esbmc_output_type,
-                    esbmc_output=esbmc_output,
-                )
-            except SourceCodeParseError:
-                pass
-            except ESBMCTimedOutException:
-                print("error: ESBMC has timed out...")
-                sys.exit(1)
-
             # Print verbose lvl 2
             print_horizontal_line(2)
             printvv(esbmc_output)
             print_horizontal_line(2)
 
+            # Solution found
             if exit_code == 0:
                 self.on_solution_signal.emit(llm_solution)
 
                 if config.raw_conversation:
                     print_raw_conversation()
 
+                printv("ESBMC-AI Notice: Succesfully verified code")
+
                 return False, llm_solution
 
+            # TODO Move this process into Solution Generator since have (beginning) is done
+            # inside, and the other half is done here.
+            # Get formatted ESBMC output
+            try:
+                esbmc_output = get_esbmc_output_formatted(
+                    esbmc_output_type=config.esbmc_output_type,
+                    esbmc_output=esbmc_output,
+                )
+            except SourceCodeParseError:
+                pass
+            except ESBMCTimedOutException:
+                print("error: ESBMC has timed out...")
+                sys.exit(1)
+
             # Failure case
-            print(f"Failure {idx+1}/{max_retries}: Retrying...")
+            print(f"ESBMC-AI Notice: Failure {idx+1}/{max_retries}: Retrying...")
             # If final iteration no need to sleep.
             if idx < max_retries - 1:
 
@@ -172,4 +176,5 @@ def print_raw_conversation() -> None:
 
         if config.raw_conversation:
             print_raw_conversation()
-        return True, "Failed all attempts..."
+
+        return True, "ESBMC-AI Notice: Failed all attempts..."