github/gpt-pilot
1
1
Fork 0
mirror of https://github.com/Pythagora-io/gpt-pilot.git synced 2026-01-09 13:17:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: Integrate NO_COLOR support and replace existing color functions

Browse Source 
- Implement NO_COLOR environment variable adherence to disable colored output.
- Replace existing output color functions with the new `color_text`.
This commit is contained in:
Umpire2018
2023-10-10 15:55:47 +08:00
parent a296d0d971
commit 5143568857
18 changed files with 129 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -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`

3
pilot/.env.example
View File

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

8
pilot/database/database.py
View File

@@ -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()

11
pilot/helpers/AgentConvo.py
View File

@@ -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")

9
pilot/helpers/Project.py
View File

@@ -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()

4
pilot/helpers/agents/Architect.py
View File

@@ -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)

36
pilot/helpers/agents/Developer.py
View File

@@ -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()

8
pilot/helpers/agents/ProductOwner.py
View File

@@ -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',

4
pilot/helpers/agents/TechLead.py
View File

@@ -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

22
pilot/helpers/cli.py
View File

@@ -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:

4
pilot/helpers/files.py
View File

@@ -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 = []

6
pilot/main.py
View File

@@ -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)

7
pilot/prompts/prompts.py
View File

@@ -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

22
pilot/utils/arguments.py
View File

@@ -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

8
pilot/utils/llm_connection.py
View File

@@ -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')

6
pilot/utils/questionary.py
View File

@@ -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:

99
pilot/utils/style.py
View File

@@ -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))

7
pilot/utils/utils.py
View File

@@ -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)