# The ChatML format
# Takes messages in and provides a text body to autocomplete on

def format(context, query, for_completion=False):

    # ChatML format:
    user_id             = "user"
    assistant_id        = "assistant"
    context_declaration = "<|im_start|>system\n"
    message_declaration = "<|im_start|>{author}\n"
    end_of_message      = "<|im_end|>\n"
    output              = ""

    if isinstance(query, str):
        query = [{"author": "user", "body": query}]

    if isinstance(query, list):
        for message in query:
            author  = message["author"]
            body    = message["body"]

            if "nickname" in message.keys():
                nickname    = message["nickname"]
                author      = nickname

            output = f"{output}{message_declaration.format(author=author)}{body}{end_of_message}"

    append = ""
    if for_completion:
        append = message_declaration.format(author=assistant_id)

    output = f"""{context_declaration}{context}{end_of_message}{output}{append}"""
    return output