import gradio as gr
from tokencost import count_message_tokens, count_string_tokens, calculate_prompt_cost, calculate_completion_cost


def compute_all(prompt_string, completion_string, model):
    prompt_cost = calculate_prompt_cost(prompt_string, model)
    completion_cost = calculate_completion_cost(completion_string, model)

    prompt_tokens = count_string_tokens(prompt_string, model)
    completion_tokens = count_string_tokens(completion_string, model)

    return prompt_tokens, prompt_cost, completion_tokens, completion_cost


with gr.Blocks(theme='soft') as demo:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(value="Hello world")
            completion = gr.Textbox(value="...")

            model = gr.Dropdown(value="gpt-3.5-turbo", choices=["gpt-3.5-turbo", "gpt-3.5-turbo"])

            button = gr.Button("Compute costs!")

        with gr.Column():
            with gr.Row():
                prompt_tokens = gr.Textbox(interactive=False)
                prompt_cost = gr.Textbox(interactive=False)
            with gr.Row():
                completion_tokens = gr.Textbox(interactive=False)
                completion_cost = gr.Textbox(interactive=False)

    button.click(compute_all, inputs=[prompt, completion, model], outputs=[prompt_tokens, prompt_cost, completion_tokens, completion_cost])

if __name__ == "__main__":
    demo.launch()