Update app.py

app.py

@@ -7,10 +7,10 @@ import gradio as gr
 from huggingface_hub import snapshot_download
 from tooluniverse import ToolUniverse
 import time
-from functools import partial
 from requests.adapters import HTTPAdapter
 from requests import Session
 from urllib3.util.retry import Retry
+from tqdm import tqdm
 
 # Configuration
 CONFIG = {
@@ -21,30 +21,42 @@ CONFIG = {
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     },
-    "download_timeout": 300,  # 5 minutes timeout
-    "max_retries": 3,
-    "retry_delay": 10  # seconds between retries
+    "download_settings": {
+        "timeout": 600,  # 10 minutes per request
+        "max_retries": 5,
+        "retry_delay": 30,  # seconds between retries
+        "chunk_size": 1024 * 1024 * 10,  # 10MB chunks
+        "max_concurrent": 2  # concurrent downloads
+    }
 }
 
 # Logging setup
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
-def create_custom_session():
-    """Create a session with custom timeout and retry settings"""
+def create_optimized_session():
+    """Create a session optimized for large file downloads"""
     session = Session()
-    retries = Retry(
-        total=CONFIG["max_retries"],
+    
+    retry_strategy = Retry(
+        total=CONFIG["download_settings"]["max_retries"],
         backoff_factor=1,
-        status_forcelist=[500, 502, 503, 504]
+        status_forcelist=[408, 429, 500, 502, 503, 504]
     )
+    
     adapter = HTTPAdapter(
-        max_retries=retries,
+        max_retries=retry_strategy,
         pool_connections=10,
-        pool_maxsize=10
+        pool_maxsize=10,
+        pool_block=True
     )
-    session.mount("http://", adapter)
+    
     session.mount("https://", adapter)
+    session.mount("http://", adapter)
+    
     return session
 
 def prepare_tool_files():
@@ -57,49 +69,66 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
-def download_with_retry(repo_id, local_dir):
-    retry_count = 0
-    custom_session = create_custom_session()
+def download_model_with_progress(repo_id, local_dir):
+    custom_session = create_optimized_session()
     
-    while retry_count < CONFIG["max_retries"]:
+    for attempt in range(CONFIG["download_settings"]["max_retries"] + 1):
         try:
+            logger.info(f"Download attempt {attempt + 1} for {repo_id}")
+            
+            # Create progress bar
+            progress = tqdm(
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                miniters=1,
+                desc=f"Downloading {repo_id.split('/')[-1]}"
+            )
+            
+            def update_progress(monitor):
+                progress.update(monitor.bytes_read - progress.n)
+            
             snapshot_download(
                 repo_id=repo_id,
                 local_dir=local_dir,
                 resume_download=True,
                 local_dir_use_symlinks=False,
                 use_auth_token=True,
+                max_workers=CONFIG["download_settings"]["max_concurrent"],
+                tqdm_class=None,  # We handle progress ourselves
                 session=custom_session
             )
+            
+            progress.close()
             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 = CONFIG["retry_delay"] * retry_count
+            logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
+            if attempt < CONFIG["download_settings"]["max_retries"]:
+                wait_time = CONFIG["download_settings"]["retry_delay"] * (attempt + 1)
                 logger.info(f"Waiting {wait_time} seconds before retry...")
                 time.sleep(wait_time)
-    return False
+            else:
+                progress.close()
+                return False
 
 def download_model_files():
     os.makedirs(CONFIG["local_dir"], exist_ok=True)
-    logger.info("Downloading model files...")
+    logger.info("Starting model downloads...")
 
     # Download main model
-    logger.info(f"Downloading {CONFIG['model_name']}...")
-    if not download_with_retry(
+    if not download_model_with_progress(
         CONFIG["model_name"],
         os.path.join(CONFIG["local_dir"], CONFIG["model_name"])
     ):
-        raise RuntimeError(f"Failed to download {CONFIG['model_name']} after {CONFIG['max_retries']} attempts")
+        raise RuntimeError(f"Failed to download {CONFIG['model_name']}")
 
     # Download RAG model
-    logger.info(f"Downloading {CONFIG['rag_model_name']}...")
-    if not download_with_retry(
+    if not download_model_with_progress(
         CONFIG["rag_model_name"],
         os.path.join(CONFIG["local_dir"], CONFIG["rag_model_name"])
     ):
-        raise RuntimeError(f"Failed to download {CONFIG['rag_model_name']} after {CONFIG['max_retries']} attempts")
+        raise RuntimeError(f"Failed to download {CONFIG['rag_model_name']}")
 
     logger.info("All model files downloaded successfully")
 
@@ -113,8 +142,7 @@ def load_embeddings(agent):
             return
         except Exception as e:
             logger.error(f"Failed to load embeddings: {e}")
-            # Fall through to generate new embeddings
-
+    
     logger.info("Generating tool embeddings...")
     try:
         tools = agent.tooluniverse.get_all_tools()
@@ -134,26 +162,35 @@ class TxAgentApp:
 
     def initialize(self):
         if self.is_initialized:
-            return "Already initialized"
-
+            return "✅ Already initialized"
+        
         try:
-            logger.info("Initializing TxAgent...")
-            self.agent = TxAgent(
-                CONFIG["model_name"],
-                CONFIG["rag_model_name"],
-                tool_files_dict=CONFIG["tool_files"],
-                force_finish=True,
-                enable_checker=True,
-                step_rag_num=10,
-                seed=100,
-                additional_default_tools=["DirectResponse", "RequireClarification"]
-            )
-            logger.info("Initializing models...")
-            self.agent.init_model()
-            logger.info("Loading embeddings...")
-            load_embeddings(self.agent)
-            self.is_initialized = True
-            logger.info("✅ TxAgent initialized successfully")
+            # Initialize with progress tracking
+            with tqdm(total=4, desc="Initializing TxAgent") as pbar:
+                logger.info("Creating TxAgent instance...")
+                self.agent = TxAgent(
+                    CONFIG["model_name"],
+                    CONFIG["rag_model_name"],
+                    tool_files_dict=CONFIG["tool_files"],
+                    force_finish=True,
+                    enable_checker=True,
+                    step_rag_num=10,
+                    seed=100,
+                    additional_default_tools=["DirectResponse", "RequireClarification"]
+                )
+                pbar.update(1)
+                
+                logger.info("Initializing models...")
+                self.agent.init_model()
+                pbar.update(1)
+                
+                logger.info("Loading embeddings...")
+                load_embeddings(self.agent)
+                pbar.update(1)
+                
+                self.is_initialized = True
+                pbar.update(1)
+                
             return "✅ TxAgent initialized successfully"
         except Exception as e:
             logger.error(f"Initialization failed: {str(e)}")
@@ -161,8 +198,8 @@ class TxAgentApp:
 
     def chat(self, message, history):
         if not self.is_initialized:
-            return history + [(message, "⚠️ Error: Model not initialized. Please click 'Initialize Model' first.")]
-
+            return history + [(message, "⚠️ Please initialize the model first")]
+        
         try:
             response = ""
             for chunk in self.agent.run_gradio_chat(
@@ -176,77 +213,85 @@ class TxAgentApp:
                 max_round=30
             ):
                 response += chunk
-
-            return history + [(message, response)]
+                yield history + [(message, response)]
+                
         except Exception as e:
             logger.error(f"Chat error: {str(e)}")
-            return history + [(message, f"Error: {str(e)}")]
+            yield history + [(message, f"Error: {str(e)}")]
 
 def create_interface():
     app = TxAgentApp()
-    with gr.Blocks(title="TxAgent", css=".gradio-container {max-width: 900px !important}") as demo:
+    
+    with gr.Blocks(
+        title="TxAgent",
+        css="""
+        .gradio-container {max-width: 900px !important}
+        .progress-bar {height: 20px !important}
+        """
+    ) as demo:
         gr.Markdown("""
-        # 🧠 TxAgent: Therapeutic Reasoning AI
-        ### A specialized AI for clinical decision support and therapeutic reasoning
+        # � TxAgent: Therapeutic Reasoning AI
+        ### Specialized for clinical decision support
         """)
-
+        
+        # Initialization section
         with gr.Row():
             init_btn = gr.Button("Initialize Model", variant="primary")
-            init_status = gr.Textbox(label="Initialization Status", interactive=False)
-
-        with gr.Row():
-            with gr.Column(scale=3):
-                chatbot = gr.Chatbot(height=600, label="Conversation", bubble_full_width=False)
-                msg = gr.Textbox(label="Your Question", placeholder="Enter your clinical question here...")
-                submit_btn = gr.Button("Submit", variant="primary")
-            with gr.Column(scale=1):
-                gr.Markdown("### Example Questions:")
-                gr.Examples(
-                    examples=[
-                        "How to adjust Journavx dosage for hepatic impairment?",
-                        "Is Xolremdi safe with Prozac for WHIM syndrome?",
-                        "Warfarin-Amiodarone contraindications?",
-                        "Alternative treatments for EGFR-positive NSCLC?"
-                    ],
-                    inputs=msg,
-                    label="Click to try"
-                )
-
+            init_status = gr.Textbox(label="Status", interactive=False)
+            download_progress = gr.Textbox(visible=False)
+        
+        # Chat interface
+        chatbot = gr.Chatbot(height=500, label="Conversation")
+        msg = gr.Textbox(label="Your clinical question", placeholder="Ask about drug interactions, dosing, etc...")
+        clear_btn = gr.Button("Clear Chat")
+        
+        # Examples
+        gr.Examples(
+            examples=[
+                "How to adjust Journavx for renal impairment?",
+                "Xolremdi and Prozac interaction in WHIM syndrome?",
+                "Alternative to Warfarin for patient with amiodarone?"
+            ],
+            inputs=msg,
+            label="Example Questions"
+        )
+        
+        # Event handlers
         init_btn.click(
             fn=app.initialize,
-            outputs=init_status,
-            api_name="initialize"
+            outputs=init_status
         )
+        
         msg.submit(
-            fn=app.chat,
-            inputs=[msg, chatbot],
-            outputs=chatbot,
-            api_name="chat"
-        )
-        submit_btn.click(
             fn=app.chat,
             inputs=[msg, chatbot],
             outputs=chatbot
         )
-
+        
+        clear_btn.click(
+            fn=lambda: ([], ""),
+            outputs=[chatbot, msg]
+        )
+    
     return demo
 
 if __name__ == "__main__":
     try:
-        logger.info("Preparing tool files...")
+        logger.info("Starting application setup...")
+        
+        # Prepare files
         prepare_tool_files()
         
-        logger.info("Downloading model files (if needed)...")
+        # Download models with progress tracking
         download_model_files()
         
-        logger.info("Launching interface...")
+        # Launch interface
         interface = create_interface()
         interface.launch(
             server_name="0.0.0.0",
             server_port=7860,
-            share=False,
-            show_error=True
+            share=False
         )
     except Exception as e:
-        logger.error(f"Application failed to start: {str(e)}")
+        logger.error(f"Fatal error: {str(e)}")
         raise