From 51435688576a4a686e088b72b9f8d7f5fdd54bdd Mon Sep 17 00:00:00 2001 From: Umpire2018 Date: Tue, 10 Oct 2023 15:55:47 +0800 Subject: [PATCH] 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 | 2 + pilot/.env.example | 3 + pilot/database/database.py | 8 +-- pilot/helpers/AgentConvo.py | 11 ++-- pilot/helpers/Project.py | 9 +-- pilot/helpers/agents/Architect.py | 4 +- pilot/helpers/agents/Developer.py | 36 +++++----- pilot/helpers/agents/ProductOwner.py | 8 +-- pilot/helpers/agents/TechLead.py | 4 +- pilot/helpers/cli.py | 22 ++++--- pilot/helpers/files.py | 4 +- pilot/main.py | 6 +- pilot/prompts/prompts.py | 7 +- pilot/utils/arguments.py | 22 +++++-- pilot/utils/llm_connection.py | 8 +-- pilot/utils/questionary.py | 6 +- pilot/utils/style.py | 99 ++++++++++------------------ pilot/utils/utils.py | 7 +- 18 files changed, 129 insertions(+), 137 deletions(-) diff --git a/README.md b/README.md index ff6f9e8c..c9fd4efc 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,8 @@ Erase all development steps previously done and continue working on an existing python main.py app_id= skip_until_dev_step=0 ``` +## `--no-color` +Disable terminal color output ## `delete_unrelated_steps` diff --git a/pilot/.env.example b/pilot/.env.example index db2ecd2c..dd8bd2a2 100644 --- a/pilot/.env.example +++ b/pilot/.env.example @@ -23,3 +23,6 @@ DB_HOST= DB_PORT= DB_USER= DB_PASSWORD= + +# Terminal has color or not +NO_COLOR= \ No newline at end of file diff --git a/pilot/database/database.py b/pilot/database/database.py index 2a19a2c5..ed6b753d 100644 --- a/pilot/database/database.py +++ b/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() diff --git a/pilot/helpers/AgentConvo.py b/pilot/helpers/AgentConvo.py index 16397851..c71cd0c5 100644 --- a/pilot/helpers/AgentConvo.py +++ b/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") diff --git a/pilot/helpers/Project.py b/pilot/helpers/Project.py index a416e3bb..7599dace 100644 --- a/pilot/helpers/Project.py +++ b/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() diff --git a/pilot/helpers/agents/Architect.py b/pilot/helpers/agents/Architect.py index c3ed8eae..fbbe4152 100644 --- a/pilot/helpers/agents/Architect.py +++ b/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) diff --git a/pilot/helpers/agents/Developer.py b/pilot/helpers/agents/Developer.py index b8dd4723..e0f3120f 100644 --- a/pilot/helpers/agents/Developer.py +++ b/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(red_bold(f'Can I retry implementing the entire development step?')) + print(color_text(f'\n--------- LLM Reached Token Limit ----------', ColorName.RED, bold=True)) + 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') + - white_bold(step['command']['command']) + - red_bold(f'\n\nCan you please make it work?')) + help_description = (color_text(f'I tried running the following command but it doesn\'t seem to work:\n\n', ColorName.RED, bold=True) + + color_text(step['command']['command'], ColorName.WHITE, bold=True) + + 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() diff --git a/pilot/helpers/agents/ProductOwner.py b/pilot/helpers/agents/ProductOwner.py index c8d0d43c..92238729 100644 --- a/pilot/helpers/agents/ProductOwner.py +++ b/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', diff --git a/pilot/helpers/agents/TechLead.py b/pilot/helpers/agents/TechLead.py index 64237187..62ddf7ea 100644 --- a/pilot/helpers/agents/TechLead.py +++ b/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 diff --git a/pilot/helpers/cli.py b/pilot/helpers/cli.py index c2f72f0c..5905e7f4 100644 --- a/pilot/helpers/cli.py +++ b/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 ----------')) - question = f'Can I execute the command: `{yellow_bold(command)}`' + print(color_text(f'\n--------- EXECUTE COMMAND ----------', ColorName.YELLOW, bold=True)) + 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: diff --git a/pilot/helpers/files.py b/pilot/helpers/files.py index a2b54cce..619ce855 100644 --- a/pilot/helpers/files.py +++ b/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 = [] diff --git a/pilot/main.py b/pilot/main.py index ceb8f4c9..79caa1e0 100644 --- a/pilot/main.py +++ b/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) diff --git a/pilot/prompts/prompts.py b/pilot/prompts/prompts.py index 371f6a3d..87a8fda8 100644 --- a/pilot/prompts/prompts.py +++ b/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 diff --git a/pilot/utils/arguments.py b/pilot/utils/arguments.py index bd305118..33dd3640 100644 --- a/pilot/utils/arguments.py +++ b/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(green_bold(f'{app.name} (app_id={arguments["app_id"]})')) - print(green_bold('--------------------------------------------------------------\n')) + print(color_text('\n------------------ LOADING PROJECT ----------------------', + ColorName.GREEN, bold=True)) + 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(green_bold('--------------------------------------------------------------\n')) + print(color_text(f'python {sys.argv[0]} app_id={arguments["app_id"]}', ColorName.GREEN, bold=True)) + 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 diff --git a/pilot/utils/llm_connection.py b/pilot/utils/llm_connection.py index 593b9804..035700f3 100644 --- a/pilot/utils/llm_connection.py +++ b/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...")) - # print(yellow("Stream response from OpenAI:")) + # spinner = spinner_start(color_text("Waiting for OpenAI API response...", ColorName.YELLOW)) + # print(color_text("Stream response from OpenAI:", ColorName.YELLOW)) # Configure for the selected ENDPOINT model = os.getenv('MODEL_NAME', 'gpt-4') diff --git a/pilot/utils/questionary.py b/pilot/utils/questionary.py index b9877b10..17b16010 100644 --- a/pilot/utils/questionary.py +++ b/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(yellow_bold(f'{user_input.user_input}')) + print(color_text(f'Restoring user input id {user_input.id}: ', ColorName.YELLOW, bold=True), end='') + 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: diff --git a/pilot/utils/style.py b/pilot/utils/style.py index 69b07366..6cd21416 100644 --- a/pilot/utils/style.py +++ b/pilot/utils/style.py @@ -1,86 +1,55 @@ -from colorama import Fore, Style, init - -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}' - - +import os +from enum import Enum from colorama import Fore, Style, init init(autoreset=True) -COLORS = { - "red": Fore.RED, - "green": Fore.GREEN, - "yellow": Fore.YELLOW, - "blue": Fore.BLUE, - "cyan": Fore.CYAN, - "white": Fore.WHITE, -} + +class ColorName(Enum): + BLACK = Fore.BLACK + RED = Fore.RED + GREEN = Fore.GREEN + YELLOW = Fore.YELLOW + 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)) diff --git a/pilot/utils/utils.py b/pilot/utils/utils.py index 03fbca81..8cd28465 100644 --- a/pilot/utils/utils.py +++ b/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)