import gradio as gr
import tiktoken
from decimal import getcontext

getcontext().prec = 2

def num_tokens_from_messages(messages, model):
    """Return the number of tokens used by a list of messages."""
    encoding = tiktoken.encoding_for_model(model)
    if model in {
        "gpt-3.5-turbo-1106",
        "gpt-4-1106-preview",
        "gpt-4",
        }:
        tokens_per_message = 3
        tokens_per_name = 1
    else:
        raise NotImplementedError(
            f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens."""
        )
    num_tokens = 0
    for message in messages:
        num_tokens += tokens_per_message
        for key, value in message.items():
            num_tokens += len(encoding.encode(value))
            if key == "name":
                num_tokens += tokens_per_name
    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
    return num_tokens

# https://openai.com/pricing
def calc_input_cost(tokens, model):
    """Return the input cost of tokens"""
    cost = 0
    if model == "gpt-3.5-turbo-1106":
      # input $0.001 / 1K tokens, output $0.002 / 1K tokens
      cost = 0.001 * tokens / 1000
    elif model == "gpt-4-1106-preview":
      # input $0.01 / 1K tokens, output $0.03 / 1K tokens
      cost = 0.01 * tokens / 1000
    elif model == "gpt-4":
      # input $0.03 / 1K tokens, output $0.06 / 1K tokens
      cost = 0.03 * tokens / 1000
    else:
      raise NotImplementedError(
            f"""calc_input_cost() is not implemented for model {model}."""
        )
    return float(cost)

# https://openai.com/pricing
def calc_output_cost(tokens, model="gpt-3.5-turbo-1106"):
    """Return the output cost of tokens"""
    cost = 0
    if model == "gpt-3.5-turbo-1106":
      # input $0.001 / 1K tokens, output $0.002 / 1K tokens
      cost = 0.002 * tokens / 1000
    elif model == "gpt-4-1106-preview":
      # input $0.03 / 1K tokens, output $0.03 / 1K tokens
      cost = 0.03 * tokens / 1000
    elif model == "gpt-4":
      # input $0.03 / 1K tokens, output $0.06 / 1K tokens
      cost = 0.06 * tokens / 1000
    else:
      raise NotImplementedError(
            f"""calc_input_cost() is not implemented for model {model}."""
        )
    return float(cost)

def make_input_messages(system, user):
    return [
        {
            "role": "system",
            "content": system,
        },
        {
            "role": "user",
            "content": user,
        }
    ]

def make_output_messages(output):
    return [
        {
            "role": "assistant",
            "content": output,
        }
    ]

def calc(request_count, system_prompt, user_prompt, output):
    result = ""
    input_messages = make_input_messages(system_prompt, user_prompt)
    output_messages = make_output_messages(output)
    for model in [
        "gpt-3.5-turbo-1106",
        "gpt-4-1106-preview",
        "gpt-4"
        ]:

        # example token count from the function defined above
        input_token = num_tokens_from_messages(input_messages, model)
        input_cost = calc_input_cost(input_token, model)
        output_token = num_tokens_from_messages(output_messages, model)
        output_cost = calc_output_cost(output_token, model)
        total_cost = input_cost + output_cost

        result += f"""{model}
        input token:{input_token}, output token: {output_token} tokens counted by tiktoken.
        input cost: ${format(input_cost, 'f')}, output cost: ${format(output_cost, 'f')}, total cost: ${format(total_cost, 'f')} per 1 request
        {int(request_count)} request cost is ${format(total_cost * request_count, 'f')}\n"""
    return result

iface = gr.Interface(
    fn=calc,
    inputs=[
        gr.Number(label="Request Count", value=500),
        gr.Textbox(label="System Prompt", value="System Prompt"),
        gr.Textbox(label="User Prompt", value="User Prompt"),
        gr.Textbox(label="Output", value="Output Text"),
    ],
    outputs=[ 
        gr.Textbox(label="Result")
    ],
    title="ChatGPT Token and Cost Calculator",
    description=f"Enter each form value, then click the button to see the results.\nThe results of this calculation should be considered for reference only."
)

iface.launch()