smarted cooldown

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from litellm import RateLimitError
 
 from time import time, sleep
 from agent import MyAgent
@@ -95,16 +96,20 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
                 additional_data = {"Content-Type": response.headers['Content-Type'],
                                "Content-Length": response.headers['Content-Length'],
                                 "content": response.content}
+                
+            cooldown_attempts = 2
+            while cooldown_attempts > 0:
+                try:
+                    submitted_answer = agent(question_text, additional_data=additional_data)
+                    break
+                except RateLimitError as e:
+                    print(f"cooldown due to {e}")
+                    sleep(61)
+                    cooldown_attempts -= 1
+                    if cooldown_attempts == 0:
+                        raise e
 
-            t0_inner = time()
-            submitted_answer = agent(question_text, additional_data=additional_data)
-            t1_inner = time()
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
-
-            if t1_inner - t0_inner > 60: # cooldown
-                print("cooldown...")
-                sleep(61)
-
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         except Exception as e:
              print(f"Error running agent on task {task_id}: {e}")