Update app.py

app.py

@@ -8,6 +8,7 @@ from i_search import google
 from i_search import i_search as i_s
 from datetime import datetime
 import logging
+import time
 
 now = datetime.now()
 date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
@@ -459,18 +460,24 @@ def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperatur
     """
     if message:
         # Run the model and get the response (convert generator to string)
-        response = ''.join(generate(
-            prompt=message,
-            history=history,
-            agent_name=agent_name,
-            sys_prompt=sys_prompt,
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-        ))
-        history.append((message, response))
-        return history
+        for attempt in range(3):  # Retry up to 3 times
+            try:
+                response = ''.join(generate(
+                    prompt=message,
+                    history=history,
+                    agent_name=agent_name,
+                    sys_prompt=sys_prompt,
+                    temperature=temperature,
+                    max_new_tokens=max_new_tokens,
+                    top_p=top_p,
+                    repetition_penalty=repetition_penalty,
+                ))
+                history.append((message, response))
+                return history
+            except huggingface_hub.errors.OverloadedError:
+                logging.warning(f"Model overloaded. Retrying in {attempt + 1} seconds...")
+                time.sleep(attempt + 1)  # Wait a bit longer each time
+        return history  # If all attempts fail, return the history without a response
 
     return history