Update app.py

app.py

@@ -3,22 +3,23 @@ import sys
 import random
 import gradio as gr
 
-# Add `src` directory to Python path
+# Add `src` to import path
 sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
 
 from txagent.txagent import TxAgent
 
-# ==== Environment Setup ====
+# === Environment setup ===
 current_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# ==== UI Content ====
+# === UI constants ===
 DESCRIPTION = '''
 <div>
-<h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools </h1>
+<h1 style="text-align: center;">TxAgent: An AI Agent for Therapeutic Reasoning Across a Universe of Tools</h1>
 </div>
 '''
+
 INTRO = "Precision therapeutics require multimodal adaptive models..."
 LICENSE = "DISCLAIMER: THIS WEBSITE DOES NOT PROVIDE MEDICAL ADVICE..."
 
@@ -49,7 +50,7 @@ chat_css = """
 .gr-button svg { width: 32px !important; height: 32px !important; }
 """
 
-# ==== Model Settings ====
+# === Model and tool config ===
 model_name = "mims-harvard/TxAgent-T1-Llama-3.1-8B"
 rag_model_name = "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B"
 new_tool_files = {
@@ -61,10 +62,10 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
-# Initialize agent placeholder
-agent = None
+agent = None  # global agent initialized later
+
 
-# ===== Build Gradio UI =====
+# === Create Gradio UI ===
 def create_ui():
     with gr.Blocks(css=css) as demo:
         gr.Markdown(DESCRIPTION)
@@ -85,7 +86,7 @@ def create_ui():
             show_copy_button=True
         )
 
-        # Retry logic
+        # === Retry handler (fixed)
         def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
             agent.update_parameters(seed=random.randint(0, 10000))
             new_history = history[:retry_data.index]
@@ -95,7 +96,7 @@ def create_ui():
                 temperature, max_new_tokens, max_tokens,
                 multi_agent, conversation, max_round
             )
-            if hasattr(result, "__iter__") and not isinstance(result, (str, dict, list)):
+            if hasattr(result, "__iter__") and not isinstance(result, (str, list, dict)):
                 result = list(result)
             return result
 
@@ -106,10 +107,14 @@ def create_ui():
             multi_agent, conversation_state, max_round
         )
 
-        # Main handler
+        # === Chat handler
         def handle_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-            return agent.run_gradio_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
+            result = agent.run_gradio_chat(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round)
+            if hasattr(result, "__iter__") and not isinstance(result, (str, list, dict)):
+                result = list(result)
+            return result
 
+        # === Chat Interface setup
         gr.ChatInterface(
             fn=handle_chat,
             chatbot=chatbot,
@@ -128,7 +133,8 @@ def create_ui():
         gr.Markdown(LICENSE)
     return demo
 
-# === VLLM-safe entry point ===
+
+# === HF Spaces + vLLM-safe entrypoint ===
 if __name__ == "__main__":
     agent = TxAgent(
         model_name,