Update app.py

app.py

@@ -19,18 +19,16 @@ state = {
 }
 
 @retry(stop=stop_after_attempt(3), wait=wait_fixed(2), retry=retry_if_exception_type(httpx.ReadTimeout))
-def get_token_count(client, api_name):
-    url = f"{client.base_url}{api_name}"
-    response = httpx.get(url, timeout=10.0)
-    response.raise_for_status()  # Raise an exception for HTTP errors
-    return int(response.json())
+def get_token_count(client):
+    result = client.predict(api_name="/update_token_display")
+    return int(result)
 
 def monitor_growth():
     while True:
         try:
-            curr_count1 = get_token_count(client1, "/update_token_display")
-            curr_count2 = get_token_count(client2, "/update_token_display")
-            curr_count3 = get_token_count(client3, "/update_token_display")
+            curr_count1 = get_token_count(client1)
+            curr_count2 = get_token_count(client2)
+            curr_count3 = get_token_count(client3)
             growth_speed1 = curr_count1 - state["prev_count1"]
             growth_speed2 = curr_count2 - state["prev_count2"]
             growth_speed3 = curr_count3 - state["prev_count3"]