Added exponential backoff

chat_application/main.py

@@ -1,3 +1,4 @@
+from google.api_core import exceptions
 from flask import Flask, request, render_template, redirect, url_for, session, make_response, render_template_string
 from flask_socketio import SocketIO, join_room, leave_room, send
 from pymongo import MongoClient
@@ -319,7 +320,7 @@ def get_response_delay(response):
 
 # Ask a bot for its response, store in DB, and send to client
     # Returns true if the bot passed
-def ask_bot(room_id, bot, bot_display_name, initial_prompt, instruct_prompt):
+def ask_bot(room_id, bot, bot_display_name, initial_prompt, instruct_prompt , wait_time = 1):
     # Prevents crashing if bot model did not load
     if bot is None:
         return False
@@ -352,6 +353,35 @@ def ask_bot(room_id, bot, bot_display_name, initial_prompt, instruct_prompt):
                 ),
             )
         parsed_response = response.candidates[0].content.parts[0].text.strip()
+
+    # Deal with rate limit issues
+    except exceptions.TooManyRequests:
+        print(f"429 Rate Limit Exceeded")
+        socketio.sleep(wait_time)
+        wait_time = wait_time *= 2
+
+        # Prevent Stack Overflow
+        if wait_time > 32:
+            print("Rate Limit Exceeded and Exponential Backoff Too Long")
+            print("Treating this bot's response as a pass.")
+            room_doc = rooms_collection.find_one({"_id": room_id})
+            if not room_doc or room_doc.get("ended", False):
+                return False
+            # Store the error response in the database
+            bot_message = {
+                "sender": bot_display_name,
+                "message": "ERROR in bot response - treated as a (pass)", 
+                "timestamp": datetime.utcnow()
+            }
+            rooms_collection.update_one(
+                {"_id": room_id},
+                {"$push": {"messages": bot_message}}
+            )
+            return True
+
+        ask_bot(room_id, bot, bot_display_name, initial_prompt, instruct_prompt , wait_time = wait_time)
+
+
     except Exception as e:
         print("Error in bot response: ", e)
         print("Treating this bot's response as a pass.")