diff --git a/scripts/print_chat.py b/scripts/print_chat.py new file mode 100644 index 0000000..a972be9 --- /dev/null +++ b/scripts/print_chat.py @@ -0,0 +1,53 @@ +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"] == "system": + formatted_messages.append(f"system: {message['content']}\n") + elif message["role"] == "user": + formatted_messages.append(f"user: {message['content']}\n") + elif message["role"] == "assistant" and message.get("function_call"): + formatted_messages.append(f"assistant: {message['function_call']}\n") + elif message["role"] == "assistant" and not message.get("function_call"): + formatted_messages.append(f"assistant: {message['content']}\n") + elif message["role"] == "function": + formatted_messages.append( + f"function ({message['name']}): {message['content']}\n" + ) + for formatted_message in formatted_messages: + print( + colored( + formatted_message, + role_to_color[ + messages[formatted_messages.index(formatted_message)]["role"] + ], + ) + ) + + +@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() diff --git a/scripts/rerun_edited_message_logs.py b/scripts/rerun_edited_message_logs.py index e0ff9bb..2bf4409 100644 --- a/scripts/rerun_edited_message_logs.py +++ b/scripts/rerun_edited_message_logs.py @@ -3,24 +3,22 @@ import pathlib import typer -from ..ai import AI -from ..chat_to_files import to_files +from gpt_engineer.ai import AI +from gpt_engineer.chat_to_files import to_files app = typer.Typer() @app.command() -def chat( +def main( messages_path: str, out_path: str | None = None, model: str = "gpt-4", temperature: float = 0.1, - max_tokens: int = 4096, ): ai = AI( model=model, temperature=temperature, - max_tokens=max_tokens, ) with open(messages_path) as f: