app.py

import gradio as gr
import pandas as pd
import requests
import json
import tiktoken

PRICES_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"

# Ensure TOKEN_COSTS is up to date when the module is loaded
try:
    response = requests.get(PRICES_URL)
    if response.status_code == 200:
        TOKEN_COSTS = response.json()
    else:
        raise Exception(f"Failed to fetch token costs, status code: {response.status_code}")
except Exception as e:
    print(f'Failed to update token costs with error: {e}. Using static costs.')
    with open("model_prices.json", "r") as f:
        TOKEN_COSTS = json.load(f)

TOKEN_COSTS = pd.DataFrame.from_dict(TOKEN_COSTS, orient='index').reset_index()
TOKEN_COSTS.columns = ['model'] + list(TOKEN_COSTS.columns[1:])
TOKEN_COSTS = TOKEN_COSTS.loc[~TOKEN_COSTS["model"].str.contains("sample_spec")]
TOKEN_COSTS = TOKEN_COSTS.loc[~TOKEN_COSTS["input_cost_per_token"].isnull()]


def count_string_tokens(string: str, model: str) -> int:
    try:
        encoding = tiktoken.encoding_for_model(model.split('/')[-1])
    except KeyError:
        print(f"Model {model} not found. Using cl100k_base encoding.")
        encoding = tiktoken.get_encoding("cl100k_base")
    return len(encoding.encode(string))

def calculate_total_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
    model_data = TOKEN_COSTS[TOKEN_COSTS['model'] == model].iloc[0]
    prompt_cost = prompt_tokens * model_data['input_cost_per_token']
    completion_cost = completion_tokens * model_data['output_cost_per_token']
    return prompt_cost, completion_cost

def update_model_list(function_calling, litellm_provider, max_price, supports_vision):
    filtered_models = TOKEN_COSTS

    if litellm_provider != "Any":
        filtered_models = filtered_models[filtered_models['litellm_provider'] == litellm_provider]
    
    if supports_vision:
        filtered_models = filtered_models[filtered_models['supports_vision']]
    
    list_models = filtered_models['model'].tolist()
    return gr.Dropdown(choices=list_models, value=list_models[0] if list_models else "No model found for this combination!")

def compute_all(input_type, prompt_text, completion_text, prompt_tokens, completion_tokens, model):
    if input_type == "Text Input":
        prompt_tokens = count_string_tokens(prompt_text, model)
        completion_tokens = count_string_tokens(completion_text, model)
    else:  # Token Count Input
        prompt_tokens = int(prompt_tokens * 1000)
        completion_tokens = int(completion_tokens * 1000)

    prompt_cost, completion_cost = calculate_total_cost(prompt_tokens, completion_tokens, model)
    total_cost = prompt_cost + completion_cost
    
    return (
        f"${prompt_cost:.6f}",
        f"${completion_cost:.6f}",
        f"${total_cost:.6f}"
    )

with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, secondary_hue=gr.themes.colors.orange)) as demo:
    gr.Markdown("""
    # Text-to-$$$: Calculate the price of your LLM runs
    Based on data from [litellm](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json).
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("## Input type")
            input_type = gr.Radio(["Text Input", "Token Count Input"], label="Input Type", value="Text Input")
            
            with gr.Group() as text_input_group:
                prompt_text = gr.Textbox(label="Prompt", value="Tell me a joke about AI. Here's an example: Why did the neural network go to therapy? It had too many deep issues!", lines=3)
                completion_text = gr.Textbox(label="Completion", value="", lines=3)
            
            with gr.Group(visible=False) as token_input_group:
                prompt_tokens_input = gr.Number(label="Prompt Tokens (thousands)", value=1.5)
                completion_tokens_input = gr.Number(label="Completion Tokens (thousands)", value=2)

            gr.Markdown("## Model choice:")
            with gr.Row():
                with gr.Column():
                    function_calling = gr.Checkbox(label="Supports Tool Calling", value=False)
                    supports_vision = gr.Checkbox(label="Supports Vision", value=False)
                litellm_provider = gr.Dropdown(label="LiteLLM Provider", choices=["Any"] + TOKEN_COSTS['litellm_provider'].unique().tolist(), value="Any")
            
            max_price = gr.Slider(label="Max Price per Token (input + output)", minimum=0, maximum=0.001, step=0.00001, value=0.001)
            
            model = gr.Dropdown(label="Model", choices=TOKEN_COSTS['model'].tolist(), value=TOKEN_COSTS['model'].tolist()[0])
            
            compute_button = gr.Button("Compute Costs", variant="secondary")

        with gr.Column(scale=1):
            with gr.Group():
                prompt_cost = gr.Textbox(label="Prompt Cost", interactive=False)
                completion_cost = gr.Textbox(label="Completion Cost", interactive=False)
                total_cost = gr.Textbox(label="Total Cost", interactive=False)

    def toggle_input_visibility(choice):
        return (
            gr.Group(visible=(choice == "Text Input")),
            gr.Group(visible=(choice == "Token Count Input"))
        )

    input_type.change(
        toggle_input_visibility,
        inputs=[input_type],
        outputs=[text_input_group, token_input_group]
    )

    # Update model list based on criteria
    function_calling.change(update_model_list, inputs=[function_calling, litellm_provider, max_price, supports_vision], outputs=model)
    litellm_provider.change(update_model_list, inputs=[function_calling, litellm_provider, max_price, supports_vision], outputs=model)
    max_price.change(update_model_list, inputs=[function_calling, litellm_provider, max_price, supports_vision], outputs=model)
    supports_vision.change(update_model_list, inputs=[function_calling, litellm_provider, max_price, supports_vision], outputs=model)

    # Compute costs
    compute_button.click(
        compute_all,
        inputs=[
            input_type,
            prompt_text,
            completion_text,
            prompt_tokens_input,
            completion_tokens_input,
            model
        ],
        outputs=[prompt_cost, completion_cost, total_cost]
    )

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