Update app.py

app.py

@@ -15,7 +15,7 @@ import aiohttp
 import asyncio
 import json
 from agent import MagAgent  
-from token_bucket import TokenBucket
+from token_bucket import Limiter
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -26,6 +26,9 @@ RATE_LIMIT = 15  # Requests per minute
 TOKEN_BUCKET_CAPACITY = RATE_LIMIT
 TOKEN_BUCKET_REFILL_RATE = RATE_LIMIT / 60.0  # Tokens per second
 
+# Initialize global token bucket
+token_bucket = Limiter(rate=TOKEN_BUCKET_REFILL_RATE, capacity=TOKEN_BUCKET_CAPACITY)
+
 async def fetch_questions(session: aiohttp.ClientSession, questions_url: str) -> list:
     """Fetch questions asynchronously."""
     try:
@@ -66,7 +69,7 @@ async def process_question(agent, question_text: str, task_id: str, results_log:
     """Process a single question with global rate limiting."""
     try:
         # Wait for a token before proceeding
-        await token_bucket.consume(1)
+        await token_bucket.wait(1)
         submitted_answer = await agent(question_text)
         results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
         return {"task_id": task_id, "submitted_answer": submitted_answer}
@@ -74,7 +77,7 @@ async def process_question(agent, question_text: str, task_id: str, results_log:
         if e.status == 429:
             print(f"Rate limit hit for task {task_id}. Retrying after delay...")
             await asyncio.sleep(60 / RATE_LIMIT)  # Wait before retry
-            await token_bucket.consume(1)
+            await token_bucket.wait(1)
             submitted_answer = await agent(question_text)
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
             return {"task_id": task_id, "submitted_answer": submitted_answer}