Update app.py

app.py

@@ -19,6 +19,7 @@ state = {
 }
 
 state_lock = Lock()
+stop_threads = False
 
 @retry(stop=stop_after_attempt(3), wait=wait_fixed(2), retry=retry_if_exception_type(httpx.ReadTimeout))
 def get_token_count(client):
@@ -26,7 +27,7 @@ def get_token_count(client):
     return int(result)
 
 def monitor_growth():
-    while True:
+    while not stop_threads:
         try:
             curr_count1 = get_token_count(client1)
             curr_count2 = get_token_count(client2)
@@ -52,7 +53,8 @@ def monitor_growth():
             print(f"An error occurred: {e}")
 
 # Start the monitor thread
-Thread(target=monitor_growth, daemon=True).start()
+monitor_thread = Thread(target=monitor_growth, daemon=True)
+monitor_thread.start()
 
 # Streamlit UI
 st.title("Token Count Dashboard")
@@ -61,10 +63,19 @@ token_count_placeholder = st.empty()
 growth_speed_placeholder = st.empty()
 
 def update_ui():
-    while True:
+    while not stop_threads:
         with state_lock:
             token_count_placeholder.metric("Total Token Count", state["prev_total_tokens"])
             growth_speed_placeholder.metric("Total Growth Speed", state["total_growth_speed"])
         time.sleep(1)
 
-Thread(target=update_ui, daemon=True).start()
+update_thread = Thread(target=update_ui, daemon=True)
+update_thread.start()
+
+try:
+    while True:
+        time.sleep(1)
+except KeyboardInterrupt:
+    stop_threads = True
+    monitor_thread.join()
+    update_thread.join()