Update app.py

app.py

@@ -1,16 +1,16 @@
-
-
 import gradio as gr
 import os
 from openai import OpenAI
+import json
+import time
 
 # OpenRouter API key
 OPENROUTER_API_KEY = "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"
 
-# Available models
+# Available models categorized by type
 TEXT_MODELS = {
     "Mistral Small": "mistralai/mistral-small-3.2-24b-instruct:free",
-    "Claude 3 Haiku": "anthropic/claude-3-haiku:free",
+    "Claude 3 Haiku": "anthropic/claude-3-haiku",
     "Gemini Pro": "google/gemini-2.5-pro-exp-03-25",
     "Qwen": "qwen/qwen2.5-32b-instruct:free",
     "Mistral 3.1": "mistralai/mistral-small-3.1-24b-instruct:free",
@@ -21,19 +21,23 @@ TEXT_MODELS = {
 # Available image models
 IMAGE_MODELS = {
     "Kimi Vision": "moonshotai/kimi-vl-a3b-thinking:free",
-    "Claude 3 Opus Vision": "anthropic/claude-3-opus-vision:free",
-    "Claude 3 Sonnet Vision": "anthropic/claude-3-sonnet-vision:free",
-    "Gemini Pro Vision": "google/gemini-pro-vision:free",
-    "GPT-4 Vision": "openai/gpt-4-vision:free",
-    "Llava": "llava/llava-1.6-34b-vision:free",
+    "Gemini Pro Vision": "google/gemini-2.5-pro-exp-03-25",
+    "Qwen Vision": "qwen/qwen2.5-vl-32b-instruct:free",
+    "Gemma Vision": "google/gemma-3-4b-it:free",
+    "Llama 3 Vision": "meta-llama/llama-3.2-11b-vision-instruct:free",
 }
 
-# Combined models for dropdown
-MODELS = {**TEXT_MODELS, **IMAGE_MODELS}
+# Token usage tracking
+token_usage = {
+    "total_input_tokens": 0,
+    "total_output_tokens": 0,
+    "model_usage": {}
+}
 
 # Initialize chat history
 history = []
 
+# Helper function to convert image to base64
 import base64
 from PIL import Image
 import io
@@ -90,15 +94,45 @@ def chat_with_ai(message, model_name, history, image=None):
             # Add the current message (text only)
             messages.append({"role": "user", "content": message})
         
+        # Get the model ID based on the selected model name
+        model_id = TEXT_MODELS.get(model_name) or IMAGE_MODELS.get(model_name)
+        
         # Create the completion request
+        start_time = time.time()
         completion = client.chat.completions.create(
             extra_headers={
                 "HTTP-Referer": "https://gradio-openrouter-interface.com",
                 "X-Title": "Gradio OpenRouter Interface",
             },
-            model=MODELS[model_name],
+            model=model_id,
             messages=messages
         )
+        end_time = time.time()
+        
+        # Update token usage statistics
+        input_tokens = completion.usage.prompt_tokens
+        output_tokens = completion.usage.completion_tokens
+        
+        token_usage["total_input_tokens"] += input_tokens
+        token_usage["total_output_tokens"] += output_tokens
+        
+        if model_name not in token_usage["model_usage"]:
+            token_usage["model_usage"][model_name] = {
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "requests": 0,
+                "avg_response_time": 0
+            }
+        
+        # Update model-specific usage
+        model_stats = token_usage["model_usage"][model_name]
+        model_stats["input_tokens"] += input_tokens
+        model_stats["output_tokens"] += output_tokens
+        model_stats["requests"] += 1
+        
+        # Calculate running average of response time
+        response_time = end_time - start_time
+        model_stats["avg_response_time"] = ((model_stats["avg_response_time"] * (model_stats["requests"] - 1)) + response_time) / model_stats["requests"]
         
         # Return the model's response
         return completion.choices[0].message.content
@@ -106,6 +140,21 @@ def chat_with_ai(message, model_name, history, image=None):
     except Exception as e:
         return f"Error: {str(e)}"
 
+# Function to generate token usage report
+def generate_usage_report():
+    report = f"### Token Usage Statistics\n\n"
+    report += f"**Total Input Tokens:** {token_usage['total_input_tokens']}\n"
+    report += f"**Total Output Tokens:** {token_usage['total_output_tokens']}\n\n"
+    
+    report += "### Model-Specific Usage\n\n"
+    report += "| Model | Input Tokens | Output Tokens | Requests | Avg Response Time (s) |\n"
+    report += "|-------|--------------|---------------|----------|----------------------|\n"
+    
+    for model, stats in token_usage["model_usage"].items():
+        report += f"| {model} | {stats['input_tokens']} | {stats['output_tokens']} | {stats['requests']} | {stats['avg_response_time']:.2f} |\n"
+    
+    return report
+
 # Create the Gradio interface
 with gr.Blocks(title="OpenRouter AI Multi-Modal Interface", css="style.css") as demo:
     gr.Markdown(
@@ -206,15 +255,18 @@ with gr.Blocks(title="OpenRouter AI Multi-Modal Interface", css="style.css") as
                             """
                             ### Available Image Models
                             - **Kimi Vision**: Moonshot AI's vision-language model
-                            - **Claude 3 Opus Vision**: Anthropic's premium vision model
-                            - **Claude 3 Sonnet Vision**: Mid-tier vision model from Anthropic
                             - **Gemini Pro Vision**: Google's multimodal vision model
-                            - **GPT-4 Vision**: OpenAI's vision-enabled GPT model
-                            - **Llava**: Open-source vision-language model
+                            - **Qwen Vision**: Alibaba's vision-language model
+                            - **Gemma Vision**: Google's lightweight vision model
+                            - **Llama 3 Vision**: Meta's vision-language model
                             
                             *Note: All responses will be provided in English*
                             """
                         )
+            
+            # Token usage statistics
+            usage_stats = gr.Markdown("### Token Usage Statistics\n\nNo usage data yet.")
+            refresh_stats_btn = gr.Button("Refresh Usage Stats")
     
     # Set up the submit action
     def respond(message, chat_history, model, image, model_type):
@@ -231,16 +283,10 @@ with gr.Blocks(title="OpenRouter AI Multi-Modal Interface", css="style.css") as
         # Process with or without image
         if use_image:
             bot_message = chat_with_ai(message, model, chat_history, image)
-            # Add a note that the response is in English
-            if not bot_message.startswith("Error:"):
-                display_message = f"{bot_message}\n\n*Response provided in English*"
-            else:
-                display_message = bot_message
         else:
             bot_message = chat_with_ai(message, model, chat_history)
-            display_message = bot_message
         
-        chat_history.append((message, display_message))
+        chat_history.append((message, bot_message))
         return "", chat_history
     
     # Connect the components
@@ -260,6 +306,9 @@ with gr.Blocks(title="OpenRouter AI Multi-Modal Interface", css="style.css") as
         return None, None
     
     clear_btn.click(clear_all, None, [chatbot, image_input], queue=False)
+    
+    # Update usage statistics
+    refresh_stats_btn.click(generate_usage_report, None, usage_stats)
 
 # For Hugging Face Spaces compatibility
 if __name__ == "__main__":
@@ -267,4 +316,5 @@ if __name__ == "__main__":
     demo.launch(share=True)
 else:
     # For Hugging Face Spaces, we need to expose the app
-    app = demo.launch(share=False, show_api=False)
+    app = demo.launch(share=False, show_api=False)
+