webbyuu / gpt-engineer /scripts /print_chat.py
NahFam13's picture
z1
d26280a verified
import json
import typer
from termcolor import colored
app = typer.Typer()
def pretty_print_conversation(messages):
role_to_color = {
"system": "red",
"user": "green",
"assistant": "blue",
"function": "magenta",
}
formatted_messages = []
for message in messages:
if message["role"] == "function":
formatted_messages.append(
f"function ({message['name']}): {message['content']}\n"
)
else:
assistant_content = (
message["function_call"]
if message.get("function_call")
else message["content"]
)
role_to_message = {
"system": f"system: {message['content']}\n",
"user": f"user: {message['content']}\n",
"assistant": f"assistant: {assistant_content}\n",
}
formatted_messages.append(role_to_message[message["role"]])
for formatted_message in formatted_messages:
role = messages[formatted_messages.index(formatted_message)]["role"]
color = role_to_color[role]
print(colored(formatted_message, color))
@app.command()
def main(
messages_path: str,
):
with open(messages_path) as f:
messages = json.load(f)
pretty_print_conversation(messages)
if __name__ == "__main__":
app()