Allow spinning to be disabled (Significant-Gravitas#4329)

MayurVirkar · May 26, 2023 · 064e95b · 064e95b
1 parent e7c0d33
commit 064e95b
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 11 deletions.
diff --git a/.env.template b/.env.template
@@ -21,6 +21,10 @@
 ## EXIT_KEY - Key to exit AUTO-GPT
 # EXIT_KEY=n
 
+## PLAIN_OUTPUT - Enabeling plain output will disable spinner (Default: False)
+## Note: Spinner is used to indicate that Auto-GPT is working on something in the background
+# PLAIN_OUTPUT=False
+
 ## DISABLED_COMMAND_CATEGORIES - The list of categories of commands that are disabled. Each of the below are an option:
 ## autogpt.commands.analyze_code
 ## autogpt.commands.audio_text

diff --git a/autogpt/agent/agent.py b/autogpt/agent/agent.py
@@ -126,7 +126,7 @@ def signal_handler(signum, frame):
                 )
                 break
             # Send message to AI, get response
-            with Spinner("Thinking... "):
+            with Spinner("Thinking... ", plain_output=cfg.plain_output):
                 assistant_reply = chat_with_ai(
                     self,
                     self.system_prompt,

diff --git a/autogpt/commands/file_operations.py b/autogpt/commands/file_operations.py
@@ -318,7 +318,7 @@ def download_file(url, filename, config: Config):
         directory = os.path.dirname(filename)
         os.makedirs(directory, exist_ok=True)
         message = f"{Fore.YELLOW}Downloading file from {Back.LIGHTBLUE_EX}{url}{Back.RESET}{Fore.RESET}"
-        with Spinner(message) as spinner:
+        with Spinner(message, plain_output=config.plain_output) as spinner:
             session = requests.Session()
             retry = Retry(total=3, backoff_factor=1, status_forcelist=[502, 503, 504])
             adapter = HTTPAdapter(max_retries=retry)

diff --git a/autogpt/config/config.py b/autogpt/config/config.py
@@ -30,6 +30,7 @@ def __init__(self) -> None:
 
         self.authorise_key = os.getenv("AUTHORISE_COMMAND_KEY", "y")
         self.exit_key = os.getenv("EXIT_KEY", "n")
+        self.plain_output = os.getenv("PLAIN_OUTPUT", "False") == "True"
 
         disabled_command_categories = os.getenv("DISABLED_COMMAND_CATEGORIES")
         if disabled_command_categories:

diff --git a/autogpt/spinner.py b/autogpt/spinner.py
@@ -8,13 +8,20 @@
 class Spinner:
     """A simple spinner class"""
 
-    def __init__(self, message: str = "Loading...", delay: float = 0.1) -> None:
+    def __init__(
+        self,
+        message: str = "Loading...",
+        delay: float = 0.1,
+        plain_output: bool = False,
+    ) -> None:
         """Initialize the spinner class
 
         Args:
             message (str): The message to display.
             delay (float): The delay between each spinner update.
+            plain_output (bool): Whether to display the spinner or not.
         """
+        self.plain_output = plain_output
         self.spinner = itertools.cycle(["-", "/", "|", "\\"])
         self.delay = delay
         self.message = message
@@ -23,11 +30,17 @@ def __init__(self, message: str = "Loading...", delay: float = 0.1) -> None:
 
     def spin(self) -> None:
         """Spin the spinner"""
+        if self.plain_output:
+            self.print_message()
+            return
         while self.running:
-            sys.stdout.write(f"{next(self.spinner)} {self.message}\r")
-            sys.stdout.flush()
+            self.print_message()
             time.sleep(self.delay)
-            sys.stdout.write(f"\r{' ' * (len(self.message) + 2)}\r")
+
+    def print_message(self):
+        sys.stdout.write(f"\r{' ' * (len(self.message) + 2)}\r")
+        sys.stdout.write(f"{next(self.spinner)} {self.message}\r")
+        sys.stdout.flush()
 
     def __enter__(self):
         """Start the spinner"""
@@ -57,9 +70,7 @@ def update_message(self, new_message, delay=0.1):
             new_message (str): New message to display.
             delay (float): The delay in seconds between each spinner update.
         """
-        time.sleep(delay)
-        sys.stdout.write(
-            f"\r{' ' * (len(self.message) + 2)}\r"
-        )  # Clear the current message
-        sys.stdout.flush()
+        self.delay = delay
         self.message = new_message
+        if self.plain_output:
+            self.print_message()
diff --git a/tests/integration/agent_factory.py b/tests/integration/agent_factory.py
@@ -12,11 +12,14 @@
 def agent_test_config(config: Config):
     was_continuous_mode = config.continuous_mode
     was_temperature = config.temperature
+    was_plain_output = config.plain_output
     config.set_continuous_mode(False)
     config.set_temperature(0)
+    config.plain_output = True
     yield config
     config.set_continuous_mode(was_continuous_mode)
     config.set_temperature(was_temperature)
+    config.plain_output = was_plain_output
 
 
 @pytest.fixture