Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from gradio_client import Client
-import time
 from collections import deque
 
 # Initialize clients
@@ -8,6 +7,14 @@ client1 = Client("orionai/training-data-collection_2")
 client2 = Client("orionai/training-data-collection_3")
 client3 = Client("orionai/training-data-collection")
 
+# URLs for the clients
+urls = {
+    client1: "https://orionai-training-data-collection-2.hf.space",
+    client2: "https://orionai-training-data-collection-3.hf.space",
+    client3: "https://orionai-training-data-collection.hf.space"
+}
+
+# Initialize state variables
 state = {
     "prev_count1": 0,
     "prev_count2": 0,
@@ -19,14 +26,16 @@ state = {
 token_counts = deque(maxlen=60)  # Store up to 60 seconds of token counts
 
 def get_token_count(client):
+    """Fetch the token count from a client."""
     try:
         result = client.predict(api_name="/update_token_display")
         return int(result)
     except Exception as e:
-        print(f"Error fetching token count: {e}")
+        print(f"Error fetching token count from {client.host}: {e}")
         return 0
 
 def update_dashboard():
+    """Update the token count and growth speed."""
     try:
         curr_count1 = get_token_count(client1)
         curr_count2 = get_token_count(client2)
@@ -45,15 +54,26 @@ def update_dashboard():
         state["prev_count3"] = curr_count3
         state["prev_total_tokens"] = total_tokens
         state["total_growth_speed"] = growth_speed
+        
+        counts = {
+            client1: curr_count1,
+            client2: curr_count2,
+            client3: curr_count3
+        }
+        
+        max_client = max(counts, key=counts.get)
+        max_tokens = counts[max_client]
+        max_url = urls[max_client]
 
-        return total_tokens, growth_speed
+        return total_tokens, growth_speed, max_tokens, max_url
     except Exception as e:
         print(f"Error in update dashboard: {e}")
-        return 0, 0
+        return 0, 0, 0, ""
 
 def refresh_metrics():
-    tokens, speed = update_dashboard()
-    return tokens, speed
+    """Refresh the metrics for the dashboard."""
+    tokens, speed, max_tokens, max_url = update_dashboard()
+    return tokens, speed, max_tokens, max_url
 
 # Create Gradio Interface
 with gr.Blocks() as demo:
@@ -61,8 +81,10 @@ with gr.Blocks() as demo:
     
     total_tokens = gr.Number(label="Total Token Count")
     growth_speed = gr.Number(label="Average Growth Speed (per second over the last minute)")
+    max_tokens = gr.Number(label="Max Token Count from a Single App")
+    max_url = gr.HTML(label="Link to App with Max Tokens")
     
     update_button = gr.Button("Update Stats")
-    update_button.click(fn=refresh_metrics, inputs=[], outputs=[total_tokens, growth_speed])
+    update_button.click(fn=refresh_metrics, inputs=[], outputs=[total_tokens, growth_speed, max_tokens, max_url])
     
 demo.launch()