Update model.py

model.py

@@ -21,21 +21,27 @@ class LocalTransformersModel:
 
 class WrapperLiteLLMModel(LiteLLMModel):
     def __call__(self, messages, **kwargs):
-        max_retry = 5
+        max_retry = 7  # Increased retries so it's more patient
         for attempt in range(max_retry):
             try:
                 return super().__call__(messages, **kwargs)
-            except RateLimitError as e:
-                print(f"RateLimitError (attempt {attempt+1}/{max_retry})")
-
-                # Try to extract retry time from the exception string
-                match = re.search(r'"retryDelay": ?"(\d+)s"', str(e))
-                retry_seconds = int(match.group(1)) if match else 50
-
-                print(f"Sleeping for {retry_seconds} seconds before retrying...")
-                time.sleep(retry_seconds)
-
-        raise RateLimitError(f"Rate limit exceeded after {max_retry} retries.")
+            except Exception as e:
+                error_str = str(e)
+                # Catch both Rate Limits (429) and Server Overloads (503)
+                if "503" in error_str or "UNAVAILABLE" in error_str or "429" in error_str:
+                    print(f"Google API busy (attempt {attempt+1}/{max_retry})")
+                    
+                    # Try to extract retry time from the exception string, default to 20 seconds
+                    match = re.search(r'"retryDelay": ?"(\d+)s"', error_str)
+                    retry_seconds = int(match.group(1)) if match else 20
+                    
+                    print(f"Sleeping for {retry_seconds} seconds before retrying...")
+                    time.sleep(retry_seconds)
+                else:
+                    # If it's a different error, crash normally
+                    raise e
+                    
+        raise Exception(f"API request failed after {max_retry} retries due to server traffic.")
 
 @lru_cache(maxsize=1)
 def get_lite_llm_model(model_id: str,  **kwargs) -> WrapperLiteLLMModel: