Better app, thx Claude

app.py

@@ -1,40 +1,105 @@
 import gradio as gr
 import pandas as pd
+import requests
+import json
+import tiktoken
 
-model_data = pd.read_json("model_prices.json")
+PRICES_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
 
-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)
+# 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)
 
-    return prompt_tokens, prompt_cost, completion_tokens, completion_cost
+TOKEN_COSTS = pd.DataFrame.from_dict(TOKEN_COSTS, orient='index').reset_index()
+TOKEN_COSTS.columns = ['model'] + list(TOKEN_COSTS.columns[1:])
 
+def count_string_tokens(string: str, model: str) -> int:
+    """Returns the number of tokens in a text string."""
+    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))
 
-with gr.Blocks(theme='soft') as demo:
-    gr.Markdown("""### Text-to-$: calculate the price of your LLM runs
-Based on data from the great [tokencost](https://github.com/AgentOps-AI/tokencost/blob/main/tokencost/model_prices.json).
-""")
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(value="Lorem ipsum dolor sit amet...")
-            completion = gr.Textbox(value="...")
+def calculate_total_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
+    """Calculate the total cost for a given model and number of tokens."""
+    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
 
-            model = gr.Dropdown(value="gpt-3.5-turbo", choices=["gpt-3.5-turbo", "gpt-3.5-turbo"])
+def update_model_list(function_calling, litellm_provider, max_price):
+    filtered_models = TOKEN_COSTS[
+        (TOKEN_COSTS['supports_function_calling'] == function_calling) &
+        (TOKEN_COSTS['litellm_provider'] == litellm_provider) &
+        (TOKEN_COSTS['input_cost_per_token'] + TOKEN_COSTS['output_cost_per_token'] <= max_price)
+    ]
+    return filtered_models['model'].tolist()
 
-            button = gr.Button("Compute costs!")
+def compute_all(prompt_string, completion_string, model):
+    prompt_tokens = count_string_tokens(prompt_string, model)
+    completion_tokens = count_string_tokens(completion_string, model)
+    cost = calculate_total_cost(prompt_tokens, completion_tokens, model)
+    prompt_cost = prompt_tokens * TOKEN_COSTS[TOKEN_COSTS['model'] == model]['input_cost_per_token'].values[0]
+    completion_cost = completion_tokens * TOKEN_COSTS[TOKEN_COSTS['model'] == model]['output_cost_per_token'].values[0]
+    
+    return (
+        f"{prompt_tokens} tokens",
+        f"${prompt_cost:.6f}",
+        f"{completion_tokens} tokens",
+        f"${completion_cost:.6f}",
+        f"${cost:.6f}"
+    )
 
-        with gr.Column():
-            with gr.Row():
-                prompt_tokens = gr.Textbox(interactive=False)
-                prompt_cost = gr.Textbox(interactive=False)
+with gr.Blocks(theme=gr.themes.Soft()) 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):
+            prompt = gr.Textbox(label="Prompt", value="Tell me a joke about AI.", lines=3)
+            completion = gr.Textbox(label="Completion", value="Here's a joke about AI: Why did the AI go to therapy? It had too many deep issues!", lines=3)
+            
             with gr.Row():
-                completion_tokens = gr.Textbox(interactive=False)
-                completion_cost = gr.Textbox(interactive=False)
+                function_calling = gr.Checkbox(label="Supports Function Calling")
+                litellm_provider = gr.Dropdown(label="LiteLLM Provider", choices=TOKEN_COSTS['litellm_provider'].unique().tolist())
+            
+            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())
+            
+            compute_button = gr.Button("Compute Costs", variant="primary")
+
+        with gr.Column(scale=1):
+            with gr.Group():
+                prompt_tokens = gr.Textbox(label="Prompt Tokens", interactive=False)
+                prompt_cost = gr.Textbox(label="Prompt Cost", interactive=False)
+                completion_tokens = gr.Textbox(label="Completion Tokens", interactive=False)
+                completion_cost = gr.Textbox(label="Completion Cost", interactive=False)
+                total_cost = gr.Textbox(label="Total Cost", interactive=False)
+
+    # Update model list based on criteria
+    function_calling.change(update_model_list, inputs=[function_calling, litellm_provider, max_price], outputs=model)
+    litellm_provider.change(update_model_list, inputs=[function_calling, litellm_provider, max_price], outputs=model)
+    max_price.change(update_model_list, inputs=[function_calling, litellm_provider, max_price], outputs=model)
 
-    button.click(compute_all, inputs=[prompt, completion, model], outputs=[prompt_tokens, prompt_cost, completion_tokens, completion_cost])
+    # Compute costs
+    compute_button.click(
+        compute_all,
+        inputs=[prompt, completion, model],
+        outputs=[prompt_tokens, prompt_cost, completion_tokens, completion_cost, total_cost]
+    )
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1 +1,2 @@
-pandas
+pandas
+tiktoken