Add vision support

app.py

@@ -21,6 +21,7 @@ except Exception as e:
 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:
@@ -37,15 +38,17 @@ def calculate_total_cost(prompt_tokens: int, completion_tokens: int, model: str)
     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):
-    if litellm_provider == "Any":
-        return gr.Dropdown(choices=TOKEN_COSTS['model'].tolist(), value=list_models[0]) 
-    else:
-        filtered_models = TOKEN_COSTS[
-            (TOKEN_COSTS['litellm_provider'] == litellm_provider)
-        ]
-        list_models = filtered_models['model'].tolist()
-        return gr.Dropdown(choices=list_models, value=list_models[0])
+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":
@@ -85,7 +88,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
 
             gr.Markdown("## Model choice:")
             with gr.Row():
-                function_calling = gr.Checkbox(label="Supports Function Calling", value=False)
+                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)
@@ -113,9 +118,10 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
     )
 
     # 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)
+    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(
@@ -132,5 +138,4 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
     )
 
 if __name__ == "__main__":
-    demo.launch()
-
+    demo.launch()