brand_tune/app.py

from brand_tune import prompting

import gradio as gr
import openai
import hashlib


SALT = "this is some random text 4e881d478cbeb1f6e5cb416770"
CORRECT_HASH = "222632858bf11308f0a7d41cb9fc061a51a47c98da85699706cc663b24aed55e"


def hash(password: str):
    return hashlib.sha256((password + SALT).encode()).hexdigest()


def gradio_history_to_openai_history(gradio_history: list[list[str]]):
    openai_history = [
        {
            "role": "system",
            "content": prompting.PROMPT_TEMPLATE,
        },
    ]
    if prompting.USE_EXAMPLE:
        openai_history += [
            {"role": "user", "content": prompting.EXAMPLE_INPUT},
            {"role": "assistant", "content": prompting.EXAMPLE_OUTPUT},
        ]

    for gradio_message in gradio_history:
        openai_history.append({"role": "user", "content": gradio_message[0]})
        if gradio_message[1]:
            openai_history.append({"role": "assistant", "content": gradio_message[1]})

    return openai_history


def bot(history: list[list[str]]):
    if hash(history[-1][0]) == CORRECT_HASH:
        history[-1][1] = "Correct! You can now use the chatbot."
        yield history
        return
    else:
        found = False
        for user_message, bot_message in history:
            if hash(user_message) == CORRECT_HASH:
                found = True
                break
        if not found:
            history[-1][1] = "Incorrect password. Try again."
            yield history
            return

    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=gradio_history_to_openai_history(history),
            stream=True,
        )
    except Exception as e:
        # An openai.error.RateLimitError can happen,
        # but we can also catch other exceptions just in case
        history[-1][1] = f"[ERROR] {type(e)}: {e}"
        yield history
        return

    history[-1][1] = ""
    for chunk in response:
        choice = chunk.choices[0]
        if choice.finish_reason is not None:
            break

        # The first chunk just says that the role is "assistant"
        # and doesn't have any content (text)
        if hasattr(choice.delta, "content"):
            history[-1][1] += choice.delta.content
            # print(choice.delta.content)

        yield history


with gr.Blocks() as interface:
    chatbot = gr.Chatbot(label="Brand Sheriff")
    msg = gr.Textbox(
        show_label=False,
        placeholder="Ex: caption a tiktok post about a weird ai-generated image",
    )
    with gr.Row():
        interrupt = gr.Button("Interrupt")
        clear = gr.Button("Clear")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    submit_event = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    interrupt.click(fn=None, inputs=None, outputs=None, cancels=[submit_event])
    clear.click(lambda: None, None, chatbot, queue=False)


interface.queue()