feat: Integrate NO_COLOR support and replace existing color functions

- Implement NO_COLOR environment variable adherence to disable colored output.
- Replace existing output color functions with the new `color_text`.

README.md

@@ -166,6 +166,8 @@ Erase all development steps previously done and continue working on an existing
 python main.py app_id=<ID_OF_THE_APP> skip_until_dev_step=0
 ```
 
+## `--no-color`
+Disable terminal color output
 
 ## `delete_unrelated_steps`
 

pilot/.env.example

@@ -23,3 +23,6 @@ DB_HOST=
 DB_PORT=
 DB_USER=
 DB_PASSWORD=
+
+# Terminal has color or not 
+NO_COLOR=

pilot/database/database.py

@@ -1,6 +1,6 @@
 from playhouse.shortcuts import model_to_dict
 from peewee import *
-from utils.style import yellow, red
+from utils.style import color_text, ColorName
 from functools import reduce
 import operator
 import psycopg2
@@ -267,7 +267,7 @@ def hash_and_save_step(Model, app_id, unique_data_fields, data_fields, message):
                        .execute())
 
         record = Model.get_by_id(inserted_id)
-        logger.debug(yellow(f"{message} with id {record.id}"))
+        logger.debug(color_text(f"{message} with id {record.id}", ColorName.YELLOW))
     except IntegrityError as e:
         logger.warn(f"A record with data {unique_data_fields} already exists for {Model.__name__}.")
         return None
@@ -373,7 +373,7 @@ def delete_all_subsequent_steps(project):
 
 
 def delete_subsequent_steps(Model, app, step):
-    logger.info(red(f"Deleting subsequent {Model.__name__} steps after {step.id if step is not None else None}"))
+    logger.info(color_text(f"Deleting subsequent {Model.__name__} steps after {step.id if step is not None else None}", ColorName.RED))
     subsequent_steps = Model.select().where(
         (Model.app == app) & (Model.previous_step == (step.id if step is not None else None)))
     for subsequent_step in subsequent_steps:
@@ -410,7 +410,7 @@ def delete_unconnected_steps_from(step, previous_step_field_name):
     ).order_by(DevelopmentSteps.id.desc())
 
     for unconnected_step in unconnected_steps:
-        print(red(f"Deleting unconnected {step.__class__.__name__} step {unconnected_step.id}"))
+        print(color_text(f"Deleting unconnected {step.__class__.__name__} step {unconnected_step.id}", ColorName.RED))
         unconnected_step.delete_instance()
 
 

pilot/helpers/AgentConvo.py

@@ -1,7 +1,7 @@
 import re
 import subprocess
 import uuid
-from utils.style import yellow, yellow_bold
+from utils.style import color_text, ColorName
 
 from database.database import get_saved_development_step, save_development_step, delete_all_subsequent_steps
 from helpers.exceptions.TokenLimitError import TokenLimitError
@@ -55,7 +55,7 @@ class AgentConvo:
         development_step = get_saved_development_step(self.agent.project)
         if development_step is not None and self.agent.project.skip_steps:
             # if we do, use it
-            print(yellow(f'Restoring development step with id {development_step.id}'))
+            print(color_text(f'Restoring development step with id {development_step.id}', ColorName.YELLOW))
             self.agent.project.checkpoints['last_development_step'] = development_step
             self.agent.project.restore_files(development_step.id)
             response = development_step.llm_response
@@ -131,7 +131,8 @@ class AgentConvo:
         # Continue conversation until GPT response equals END_RESPONSE
         while response != END_RESPONSE:
             user_message = ask_user(self.agent.project, response,
-                                    hint=yellow("Do you want to add anything else? If not, ") + yellow_bold('just press ENTER.'),
+                                    hint=color_text("Do you want to add anything else? If not, ", ColorName.YELLOW)
+                                         + color_text('just press ENTER.', ColorName.YELLOW, bold=True),
                                     require_some_input=False)
 
             if user_message == "":
@@ -191,7 +192,9 @@ class AgentConvo:
         print_msg = capitalize_first_word_with_underscores(self.high_level_step)
         if self.log_to_user:
             if self.agent.project.checkpoints['last_development_step'] is not None:
-                print(yellow("\nDev step ") + yellow_bold(str(self.agent.project.checkpoints['last_development_step'])) + '\n', end='')
+                print(color_text("\nDev step ")
+                      + color_text(str(self.agent.project.checkpoints['last_development_step']),
+                                   ColorName.YELLOW, bold=True) + '\n', end='')
             print(f"\n{content}\n", type='local')
         logger.info(f"{print_msg}: {content}\n")
 

pilot/helpers/Project.py

@@ -1,7 +1,7 @@
 import json
 import os
 from typing import Tuple
-from utils.style import  yellow_bold, cyan, white_bold
+from utils.style import color_text, ColorName
 from const.common import IGNORE_FOLDERS, STEPS
 from database.database import delete_unconnected_steps_from, delete_all_app_development_data
 from const.ipc import MESSAGE_TYPE
@@ -300,7 +300,7 @@ class Project:
         development_step, created = DevelopmentSteps.get_or_create(id=development_step_id)
 
         for file in files:
-            print(cyan(f'Saving file {(file["path"])}/{file["name"]}'))
+            print(color_text(f'Saving file {(file["path"])}/{file["name"]}', ColorName.CYAN))
             # TODO this can be optimized so we don't go to the db each time
             file_in_db, created = File.get_or_create(
                 app=self.app,
@@ -333,10 +333,11 @@ class Project:
 
     def ask_for_human_intervention(self, message, description=None, cbs={}, convo=None, is_root_task=False):
         answer = ''
-        question = yellow_bold(message)
+        question = color_text(message, ColorName.YELLOW, bold=True)
 
         if description is not None:
-            question += '\n' + '-' * 100 + '\n' + white_bold(description) + '\n' + '-' * 100 + '\n'
+            question += ('\n' + '-' * 100 + '\n' + color_text(description, ColorName.WHITE, bold=True)
+                         + '\n' + '-' * 100 + '\n')
 
         reset_branch_id = None if convo is None else convo.save_branch()
 

pilot/helpers/agents/Architect.py

@@ -1,7 +1,7 @@
 from utils.utils import step_already_finished
 from helpers.Agent import Agent
 import json
-from utils.style import green_bold
+from utils.style import color_text, ColorName
 from const.function_calls import ARCHITECTURE
 
 from utils.utils import should_execute_step, find_role_from_step, generate_app_data
@@ -28,7 +28,7 @@ class Architect(Agent):
             return step['architecture']
 
         # ARCHITECTURE
-        print(green_bold(f"Planning project architecture...\n"))
+        print(color_text(f"Planning project architecture...\n", ColorName.GREEN, bold=True))
         logger.info(f"Planning project architecture...")
 
         self.convo_architecture = AgentConvo(self)

pilot/helpers/agents/Developer.py

@@ -1,5 +1,5 @@
 import uuid
-from utils.style import green, red, green_bold, yellow_bold, red_bold, blue_bold, white_bold
+from utils.style import color_text, ColorName
 from helpers.exceptions.TokenLimitError import TokenLimitError
 from const.code_execution import MAX_COMMAND_DEBUG_TRIES
 from helpers.exceptions.TooDeepRecursionError import TooDeepRecursionError
@@ -33,7 +33,7 @@ class Developer(Agent):
             self.project.skip_steps = False if ('skip_until_dev_step' in self.project.args and self.project.args['skip_until_dev_step'] == '0') else True
 
         # DEVELOPMENT
-        print(green_bold(f"🚀 Now for the actual development...\n"))
+        print(color_text(f"🚀 Now for the actual development...\n", ColorName.GREEN, bold=True))
         logger.info(f"Starting to create the actual code...")
 
         for i, dev_task in enumerate(self.project.development_plan):
@@ -44,7 +44,8 @@ class Developer(Agent):
         logger.info('The app is DONE!!! Yay...you can use it now.')
 
     def implement_task(self, i, development_task=None):
-        print(green_bold(f'Implementing task #{i + 1}: ') + green(f' {development_task["description"]}\n'))
+        print(color_text(f'Implementing task #{i + 1}: ')
+              + color_text(f' {development_task["description"]}\n', ColorName.GREEN, bold=True))
 
         convo_dev_task = AgentConvo(self)
         task_description = convo_dev_task.send_message('development/task/breakdown.prompt', {
@@ -115,7 +116,7 @@ class Developer(Agent):
 
         while True:
             human_intervention_description = step['human_intervention_description'] + \
-                                             yellow_bold('\n\nIf you want to run the app, just type "r" and press ENTER and that will run `' + self.run_command + '`') \
+                                             color_text('\n\nIf you want to run the app, just type "r" and press ENTER and that will run `' + self.run_command + '`', ColorName.YELLOW, bold=True) \
                                              if self.run_command is not None else step['human_intervention_description']
             response = self.project.ask_for_human_intervention('I need human intervention:',
                 human_intervention_description,
@@ -148,9 +149,9 @@ class Developer(Agent):
             cli_response, llm_response = execute_command_and_check_cli_response(test_command['command'], test_command['timeout'], convo)
             logger.info('After running command llm_response: ' + llm_response)
             if llm_response == 'NEEDS_DEBUGGING':
-                print(red(f'Got incorrect CLI response:'))
+                print(color_text(f'Got incorrect CLI response:', ColorName.RED))
                 print(cli_response)
-                print(red('-------------------'))
+                print(color_text('-------------------', ColorName.RED))
 
             return { "success": llm_response == 'DONE', "cli_response": cli_response, "llm_response": llm_response }
 
@@ -183,8 +184,8 @@ class Developer(Agent):
         if step_implementation_try >= MAX_COMMAND_DEBUG_TRIES:
             self.dev_help_needed(step)
 
-        print(red_bold(f'\n--------- LLM Reached Token Limit ----------'))
+        print(color_text(f'\n--------- LLM Reached Token Limit ----------', ColorName.RED, bold=True))
-        print(red_bold(f'Can I retry implementing the entire development step?'))
+        print(color_text(f'Can I retry implementing the entire development step?', ColorName.RED, bold=True))
 
         answer = ''
         while answer != 'y':
@@ -202,9 +203,9 @@ class Developer(Agent):
     def dev_help_needed(self, step):
 
         if step['type'] == 'command':
-            help_description = (red_bold(f'I tried running the following command but it doesn\'t seem to work:\n\n') +
+            help_description = (color_text(f'I tried running the following command but it doesn\'t seem to work:\n\n', ColorName.RED, bold=True) +
-                white_bold(step['command']['command']) +
+                color_text(step['command']['command'], ColorName.WHITE, bold=True) +
-                red_bold(f'\n\nCan you please make it work?'))
+                color_text(f'\n\nCan you please make it work?', ColorName.RED, bold=True))
         elif step['type'] == 'code_change':
             help_description = step['code_change_description']
         elif step['type'] == 'human_intervention':
@@ -223,9 +224,11 @@ class Developer(Agent):
 
         answer = ''
         while answer != 'continue':
-            print(red_bold(f'\n----------------------------- I need your help ------------------------------'))
+            print(color_text(f'\n----------------------------- I need your help ------------------------------',
+                             ColorName.RED, bold=True))
             print(extract_substring(str(help_description)))
-            print(red_bold(f'\n-----------------------------------------------------------------------------'))
+            print(color_text(f'\n-----------------------------------------------------------------------------',
+                             ColorName.RED, bold=True))
             answer = styled_text(
                 self.project,
                 'Once you\'re done, type "continue"?'
@@ -291,11 +294,12 @@ class Developer(Agent):
         while True:
             logger.info('Continue development: %s', last_branch_name)
             iteration_convo.load_branch(last_branch_name)
-            user_description = ('Here is a description of what should be working: \n\n' + blue_bold(continue_description) + '\n') \
+            user_description = ('Here is a description of what should be working: \n\n' + color_text(continue_description
+                                + '\n', ColorName.BLUE, bold=True)) \
                                 if continue_description != '' else ''
             user_description = 'Can you check if the app works please? ' + user_description + \
                                '\nIf you want to run the app, ' + \
-                               yellow_bold('just type "r" and press ENTER and that will run `' + self.run_command + '`')
+                               color_text('just type "r" and press ENTER and that will run `' + self.run_command + '`', ColorName.YELLOW, bold=True)
             # continue_description = ''
             # TODO: Wait for a specific string in the output or timeout?
             response = self.project.ask_for_human_intervention(
@@ -355,7 +359,7 @@ class Developer(Agent):
         })
         return
         # ENVIRONMENT SETUP
-        print(green(f"Setting up the environment...\n"))
+        print(color_text(f"Setting up the environment...\n", ColorName.GREEN))
         logger.info(f"Setting up the environment...")
 
         os_info = get_os_info()

pilot/helpers/agents/ProductOwner.py

@@ -1,5 +1,5 @@
 import json
-from utils.style import green_bold
+from utils.style import color_text, ColorName
 from helpers.AgentConvo import AgentConvo
 from helpers.Agent import Agent
 from logger.logger import logger
@@ -55,7 +55,7 @@ class ProductOwner(Agent):
             self.project,
             generate_messages_from_description(main_prompt, self.project.args['app_type'], self.project.args['name']))
 
-        print(green_bold('Project Summary:\n'))
+        print(color_text('Project Summary:\n', ColorName.GREEN, bold=True))
         convo_project_description = AgentConvo(self)
         high_level_summary = convo_project_description.send_message('utils/summary.prompt',
                                                                     {'conversation': '\n'.join(
@@ -87,7 +87,7 @@ class ProductOwner(Agent):
 
         # USER STORIES
         msg = f"User Stories:\n"
-        print(green_bold(msg))
+        print(color_text(msg, ColorName.GREEN, bold=True))
         logger.info(msg)
 
         self.project.user_stories = self.convo_user_stories.continuous_conversation('user_stories/specs.prompt', {
@@ -121,7 +121,7 @@ class ProductOwner(Agent):
 
         # USER TASKS
         msg = f"User Tasks:\n"
-        print(green_bold(msg))
+        print(color_text(msg, ColorName.GREEN, bold=True))
         logger.info(msg)
 
         self.project.user_tasks = self.convo_user_stories.continuous_conversation('user_stories/user_tasks.prompt',

pilot/helpers/agents/TechLead.py

@@ -1,7 +1,7 @@
 from utils.utils import step_already_finished
 from helpers.Agent import Agent
 import json
-from utils.style import green_bold
+from utils.style import color_text, ColorName
 from const.function_calls import DEV_STEPS
 from helpers.cli import build_directory_tree
 from helpers.AgentConvo import AgentConvo
@@ -32,7 +32,7 @@ class TechLead(Agent):
             return step['development_plan']
         
         # DEVELOPMENT PLANNING
-        print(green_bold(f"Starting to create the action plan for development...\n"))
+        print(color_text(f"Starting to create the action plan for development...\n", ColorName.GREEN, bold=True))
         logger.info(f"Starting to create the action plan for development...")
 
         # TODO add clarifications

pilot/helpers/cli.py

@@ -8,7 +8,7 @@ import platform
 from typing import Dict, Union
 
 from logger.logger import logger
-from utils.style import yellow, green, red, yellow_bold, white_bold
+from utils.style import color_text, ColorName
 from database.database import get_saved_command_run, save_command_run
 from helpers.exceptions.TooDeepRecursionError import TooDeepRecursionError
 from helpers.exceptions.TokenLimitError import TokenLimitError
@@ -133,8 +133,8 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
             timeout = min(max(timeout, MIN_COMMAND_RUN_TIME), MAX_COMMAND_RUN_TIME)
 
     if not force:
-        print(yellow_bold(f'\n--------- EXECUTE COMMAND ----------'))
+        print(color_text(f'\n--------- EXECUTE COMMAND ----------', ColorName.YELLOW, bold=True))
-        question = f'Can I execute the command: `{yellow_bold(command)}`'
+        question = f'Can I execute the command: `{color_text(command, ColorName.YELLOW, bold=True)}`'
         if timeout is not None:
             question += f' with {timeout}ms timeout?'
         else:
@@ -162,7 +162,8 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
     if command_run is not None and project.skip_steps:
         # if we do, use it
         project.checkpoints['last_command_run'] = command_run
-        print(yellow(f'Restoring command run response id {command_run.id}:\n```\n{command_run.cli_response}```'))
+        print(color_text(f'Restoring command run response id {command_run.id}:\n```\n{command_run.cli_response}```',
+                         ColorName.YELLOW))
         return command_run.cli_response, None, None
 
     return_value = None
@@ -191,7 +192,8 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
             elapsed_time = time.time() - start_time
             if timeout is not None:
                 # TODO: print to IPC using a different message type so VS Code can ignore it or update the previous value
-                print(white_bold(f'\rt: {round(elapsed_time * 1000)}ms : '), end='', flush=True)
+                print(color_text(f'\rt: {round(elapsed_time * 1000)}ms : ', ColorName.WHITE, bold=True),
+                      end='', flush=True,)
 
             # Check if process has finished
             if process.poll() is not None:
@@ -200,7 +202,7 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
                 while not q.empty():
                     output_line = q.get_nowait()
                     if output_line not in output:
-                        print(green('CLI OUTPUT:') + output_line, end='')
+                        print(color_text('CLI OUTPUT:', ColorName.GREEN) + output_line, end='', )
                         logger.info('CLI OUTPUT: ' + output_line)
                         output += output_line
                 break
@@ -218,7 +220,7 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
 
             if line:
                 output += line
-                print(green('CLI OUTPUT:') + line, end='')
+                print(color_text('CLI OUTPUT:', ColorName.GREEN) + line, end='')
                 logger.info('CLI OUTPUT: ' + line)
 
             # Read stderr
@@ -229,7 +231,7 @@ def execute_command(project, command, timeout=None, process_name: str = None, fo
 
             if stderr_line:
                 stderr_output += stderr_line
-                print(red('CLI ERROR:') + stderr_line, end='')  # Print with different color for distinction
+                print(color_text('CLI ERROR:', ColorName.RED) + stderr_line, end='')  # Print with different color for distinction
                 logger.error('CLI ERROR: ' + stderr_line)
                 
             if process_name is not None:
@@ -373,9 +375,9 @@ def run_command_until_success(convo, command,
 
     if response != 'DONE':
         # 'NEEDS_DEBUGGING'
-        print(red(f'Got incorrect CLI response:'))
+        print(color_text(f'Got incorrect CLI response:', ColorName.RED))
         print(cli_response)
-        print(red('-------------------'))
+        print(color_text('-------------------', ColorName.RED))
 
         reset_branch_id = convo.save_branch()
         while True:

pilot/helpers/files.py

@@ -1,4 +1,4 @@
-from utils.style import green
+from utils.style import color_text, ColorName
 import os
 
 
@@ -11,7 +11,7 @@ def update_file(path, new_content):
     # Write content to the file
     with open(path, 'w') as file:
         file.write(new_content)
-        print(green(f"Updated file {path}"))
+        print(color_text(f"Updated file {path}", ColorName.GREEN))
 
 def get_files_content(directory, ignore=[]):
     return_array = []

pilot/main.py

@@ -8,7 +8,7 @@ import traceback
 from dotenv import load_dotenv
 load_dotenv()
 
-from utils.style import red
+from utils.style import color_text, ColorName
 from utils.custom_print import get_custom_print
 from helpers.Project import Project
 from utils.arguments import get_arguments
@@ -53,9 +53,9 @@ if __name__ == "__main__":
     except KeyboardInterrupt:
         exit_gpt_pilot()
     except Exception as e:
-        print(red('---------- GPT PILOT EXITING WITH ERROR ----------'))
+        print(color_text('---------- GPT PILOT EXITING WITH ERROR ----------', ColorName.RED))
         traceback.print_exc()
-        print(red('--------------------------------------------------'))
+        print(color_text('--------------------------------------------------', ColorName.RED))
         exit_gpt_pilot(False)
     finally:
         sys.exit(0)

pilot/prompts/prompts.py

@@ -1,5 +1,5 @@
 # prompts/prompts.py
-from utils.style import yellow
+from utils.style import color_text, ColorName
 from const import common
 from const.llm import MAX_QUESTIONS, END_RESPONSE
 from utils.llm_connection import create_gpt_chat_completion
@@ -52,7 +52,7 @@ def ask_for_main_app_definition(project):
 def ask_user(project, question: str, require_some_input=True, hint: str = None):
     while True:
         if hint is not None:
-            print(yellow(hint), type='hint')
+            print(color_text(hint, ColorName.YELLOW), type='hint')
         answer = styled_text(project, question)
 
         logger.info('Q: %s', question)
@@ -126,7 +126,8 @@ def get_additional_info_from_user(project, messages, role):
         while True:
             if isinstance(message, dict) and 'text' in message:
                 message = message['text']
-            print(yellow(f"Please check this message and say what needs to be changed. If everything is ok just press ENTER",))
+            print(color_text(f"Please check this message and say what needs to be changed. If everything is ok just press ENTER",
+                             ColorName.YELLOW))
             answer = ask_user(project, message, require_some_input=False)
             if answer.lower() == '':
                 break

pilot/utils/arguments.py

@@ -5,7 +5,7 @@ import sys
 import uuid
 from getpass import getuser
 from database.database import get_app, get_app_by_user_workspace
-from utils.style import green_bold
+from utils.style import color_text, ColorName
 from utils.utils import should_execute_step
 
 
@@ -25,6 +25,9 @@ def get_arguments():
         else:
             arguments[arg] = True
 
+    if '--no-color' in args:
+        os.environ["NO_COLOR"] = True
+
     if 'user_id' not in arguments:
         arguments['user_id'] = username_to_uuid(getuser())
 
@@ -46,18 +49,22 @@ def get_arguments():
             if 'step' not in arguments or ('step' in arguments and not should_execute_step(arguments['step'], app.status)):
                 arguments['step'] = app.status
 
-            print(green_bold('\n------------------ LOADING PROJECT ----------------------'))
+            print(color_text('\n------------------ LOADING PROJECT ----------------------',
-            print(green_bold(f'{app.name} (app_id={arguments["app_id"]})'))
+                             ColorName.GREEN, bold=True))
-            print(green_bold('--------------------------------------------------------------\n'))
+            print(color_text(f'{app.name} (app_id={arguments["app_id"]})', ColorName.GREEN, bold=True))
+            print(color_text('--------------------------------------------------------------\n',
+                             ColorName.GREEN, bold=True))
         except ValueError as e:
             print(e)
             exit(1)
     else:
         arguments['app_id'] = str(uuid.uuid4())
-        print(green_bold('\n------------------ STARTING NEW PROJECT ----------------------'))
+        print(color_text('\n------------------ STARTING NEW PROJECT ----------------------',
+                         ColorName.GREEN, bold=True))
         print("If you wish to continue with this project in future run:")
-        print(green_bold(f'python {sys.argv[0]} app_id={arguments["app_id"]}'))
+        print(color_text(f'python {sys.argv[0]} app_id={arguments["app_id"]}', ColorName.GREEN, bold=True))
-        print(green_bold('--------------------------------------------------------------\n'))
+        print(color_text('--------------------------------------------------------------\n',
+                         ColorName.GREEN, bold=True))
 
     if 'email' not in arguments:
         arguments['email'] = get_email()
@@ -68,6 +75,7 @@ def get_arguments():
     if 'step' not in arguments:
         arguments['step'] = None
 
+
     return arguments
 
 

pilot/utils/llm_connection.py

@@ -8,7 +8,7 @@ import tiktoken
 from prompt_toolkit.styles import Style
 
 from jsonschema import validate, ValidationError
-from utils.style import red
+from utils.style import color_text, ColorName
 from typing import List
 from const.llm import MIN_TOKENS_FOR_GPT_RESPONSE, MAX_GPT_MODEL_TOKENS
 from logger.logger import logger, logging
@@ -217,7 +217,7 @@ def retry_on_exception(func):
                         time.sleep(wait_duration)
                     continue
 
-                print(red(f'There was a problem with request to openai API:'))
+                print(color_text(f'There was a problem with request to openai API:', ColorName.RED))
                 # spinner_stop(spinner)
                 print(err_str)
                 logger.error(f'There was a problem with request to openai API: {err_str}')
@@ -284,8 +284,8 @@ def stream_gpt_completion(data, req_type, project):
         delete_last_n_lines(lines_printed)
         return result_data
 
-    # spinner = spinner_start(yellow("Waiting for OpenAI API response..."))
+    # spinner = spinner_start(color_text("Waiting for OpenAI API response...", ColorName.YELLOW))
-    # print(yellow("Stream response from OpenAI:"))
+    # print(color_text("Stream response from OpenAI:", ColorName.YELLOW))
 
     # Configure for the selected ENDPOINT
     model = os.getenv('MODEL_NAME', 'gpt-4')

pilot/utils/questionary.py

@@ -1,6 +1,6 @@
 from prompt_toolkit.styles import Style
 import questionary
-from utils.style import yellow_bold
+from utils.style import color_text, ColorName
 import re
 from database.database import save_user_input, get_saved_user_input
 
@@ -30,8 +30,8 @@ def styled_text(project, question, ignore_user_input_count=False, style=None):
         if user_input is not None and user_input.user_input is not None and project.skip_steps:
             # if we do, use it
             project.checkpoints['last_user_input'] = user_input
-            print(yellow_bold(f'Restoring user input id {user_input.id}: '), end='')
+            print(color_text(f'Restoring user input id {user_input.id}: ', ColorName.YELLOW, bold=True), end='')
-            print(yellow_bold(f'{user_input.user_input}'))
+            print(color_text(f'{user_input.user_input}', ColorName.YELLOW, bold=True))
             return user_input.user_input
 
     if project.ipc_client_instance is None or project.ipc_client_instance.client is None:

pilot/utils/style.py

@@ -1,86 +1,55 @@
-from colorama import Fore, Style, init
+import os
-
+from enum import Enum
-init()
-
-def red(text):
-    return f'{Fore.RED}{text}{Style.RESET_ALL}'
-
-
-def red_bold(text):
-    return f'{Fore.RED}{Style.BRIGHT}{text}{Style.RESET_ALL}'
-
-
-def yellow(text):
-    return f'{Fore.YELLOW}{text}{Style.RESET_ALL}'
-
-
-def yellow_bold(text):
-    return f'{Fore.YELLOW}{Style.BRIGHT}{text}{Style.RESET_ALL}'
-
-
-def green(text):
-    return f'{Fore.GREEN}{text}{Style.RESET_ALL}'
-
-
-def green_bold(text):
-    return f'{Fore.GREEN}{Style.BRIGHT}{text}{Style.RESET_ALL}'
-
-
-def blue(text):
-    return f'{Fore.BLUE}{text}{Style.RESET_ALL}'
-
-
-def blue_bold(text):
-    return f'{Fore.BLUE}{Style.BRIGHT}{text}{Style.RESET_ALL}'
-
-
-def cyan(text):
-    return f'{Fore.CYAN}{text}{Style.RESET_ALL}'
-
-
-def white(text):
-    return f'{Fore.WHITE}{text}{Style.RESET_ALL}'
-
-
-def white_bold(text):
-    return f'{Fore.WHITE}{Style.BRIGHT}{text}{Style.RESET_ALL}'
-
-
 from colorama import Fore, Style, init
 
 init(autoreset=True)
 
-COLORS = {
+
-    "red": Fore.RED,
+class ColorName(Enum):
-    "green": Fore.GREEN,
+    BLACK = Fore.BLACK
-    "yellow": Fore.YELLOW,
+    RED = Fore.RED
-    "blue": Fore.BLUE,
+    GREEN = Fore.GREEN
-    "cyan": Fore.CYAN,
+    YELLOW = Fore.YELLOW
-    "white": Fore.WHITE,
+    BLUE = Fore.BLUE
-}
+    CYAN = Fore.CYAN
+    WHITE = Fore.WHITE
 
 
-def color_text(text: str, color_name: str, bold: bool = False, enable_formatting: bool = True) -> str:
+class Config:
+    # Use environment variable or default value for NO_COLOR
+    NO_COLOR = os.environ.get("NO_COLOR", False)
+
+
+def color_text(text: str, color_name: ColorName, bold: bool = False) -> str:
     """
     Returns text with a specified color and optional style.
 
     Args:
         text (str): The text to colorize.
-        color_name (str): The color of the text. Should be a key in the COLORS dictionary.
+        color_name (ColorName): The color of the text. Should be a member of the ColorName enum.
         bold (bool, optional): If True, the text will be displayed in bold. Defaults to False.
-        enable_formatting (bool, optional): If True, ANSI codes will be applied. Defaults to True.
 
     Returns:
         str: The text with applied color and optional style.
-
-    Example:
-        >>> color_text("Hello, World!", "red", True, enable_formatting=False)
-        'Hello, World!'
     """
-    if not enable_formatting:
+    # If NO_COLOR is True, return unmodified text
+    if Config.NO_COLOR:
         return text
 
-    color = COLORS.get(color_name, Fore.WHITE)  # Default color is white
+    # Use the color value from the enum
+    color = color_name.value
+    # Apply BRIGHT style if bold is True, otherwise use an empty string
     style = Style.BRIGHT if bold else ""
+    # Return the formatted text
     return f'{color}{style}{text}'
 
+
+# Example usage:
+if __name__ == "__main__":
+    print(color_text("This is black text", ColorName.BLACK))
+    print(color_text("This is red text", ColorName.RED))
+    print(color_text("This is green text", ColorName.GREEN, bold=True))
+    print(color_text("This is yellow text", ColorName.YELLOW))
+    print(color_text("This is blue text", ColorName.BLUE, bold=True))
+    print(color_text("This is cyan text", ColorName.CYAN))
+    print(color_text("This is white text", ColorName.WHITE, bold=True))

pilot/utils/utils.py

@@ -9,8 +9,7 @@ import json
 import hashlib
 import re
 from jinja2 import Environment, FileSystemLoader
-from .style import green
+from .style import color_text, ColorName
-
 from const.llm import MAX_QUESTIONS, END_RESPONSE
 from const.common import ROLES, STEPS
 from logger.logger import logger
@@ -114,7 +113,7 @@ def get_os_info():
         "Node": platform.node(),
         "Release": platform.release(),
     }
-
+    # TODO deprecated linux_distribution
     if os_info["OS"] == "Linux":
         os_info["Distribution"] = ' '.join(distro.linux_distribution(full_distribution_name=True))
     elif os_info["OS"] == "Windows":
@@ -142,7 +141,7 @@ def step_already_finished(args, step):
     args.update(step['app_data'])
 
     message = f"✅  {capitalize_first_word_with_underscores(step['step'])}"
-    print(green(message))
+    print(color_text(message, ColorName.GREEN))
     logger.info(message)