Update app.py

app.py

@@ -4,10 +4,13 @@ import logging
 import torch
 from txagent import TxAgent
 import gradio as gr
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import snapshot_download
 from tooluniverse import ToolUniverse
-from tqdm import tqdm
 import time
+from functools import partial
+from requests.adapters import HTTPAdapter
+from requests import Session
+from urllib3.util.retry import Retry
 
 # Configuration
 CONFIG = {
@@ -18,14 +21,32 @@ CONFIG = {
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     },
-    "download_timeout": 300,  # Increased timeout to 5 minutes
-    "max_retries": 3
+    "download_timeout": 300,  # 5 minutes timeout
+    "max_retries": 3,
+    "retry_delay": 10  # seconds between retries
 }
 
 # Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+def create_custom_session():
+    """Create a session with custom timeout and retry settings"""
+    session = Session()
+    retries = Retry(
+        total=CONFIG["max_retries"],
+        backoff_factor=1,
+        status_forcelist=[500, 502, 503, 504]
+    )
+    adapter = HTTPAdapter(
+        max_retries=retries,
+        pool_connections=10,
+        pool_maxsize=10
+    )
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
+
 def prepare_tool_files():
     os.makedirs("./data", exist_ok=True)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
@@ -38,6 +59,8 @@ def prepare_tool_files():
 
 def download_with_retry(repo_id, local_dir):
     retry_count = 0
+    custom_session = create_custom_session()
+    
     while retry_count < CONFIG["max_retries"]:
         try:
             snapshot_download(
@@ -45,14 +68,15 @@ def download_with_retry(repo_id, local_dir):
                 local_dir=local_dir,
                 resume_download=True,
                 local_dir_use_symlinks=False,
-                timeout=CONFIG["download_timeout"]
+                use_auth_token=True,
+                session=custom_session
             )
             return True
         except Exception as e:
             retry_count += 1
             logger.error(f"Attempt {retry_count} failed for {repo_id}: {str(e)}")
             if retry_count < CONFIG["max_retries"]:
-                wait_time = 10 * retry_count
+                wait_time = CONFIG["retry_delay"] * retry_count
                 logger.info(f"Waiting {wait_time} seconds before retry...")
                 time.sleep(wait_time)
     return False