Update app.py

app.py

@@ -2,12 +2,10 @@ import os
 import sys
 import random
 import gradio as gr
-from datetime import datetime
 
 # Add `src` directory to Python path
 sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
 
-# Import your agent class from src/txagent/txagent.py
 from txagent.txagent import TxAgent
 
 # ==== Environment Setup ====
@@ -63,85 +61,86 @@ question_examples = [
     ["A 30-year-old patient is on Prozac for depression and now diagnosed with WHIM syndrome. Is Xolremdi suitable?"]
 ]
 
-# === Initialize the model ===
-agent = TxAgent(
-    model_name,
-    rag_model_name,
-    tool_files_dict=new_tool_files,
-    force_finish=True,
-    enable_checker=True,
-    step_rag_num=10,
-    seed=100,
-    additional_default_tools=["DirectResponse", "RequireClarification"]
-)
-agent.init_model()
-
-# === Gradio interface logic ===
-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)
-
-def update_seed():
-    seed = random.randint(0, 10000)
-    return agent.update_parameters(seed=seed)
-
-# ✅ FIXED: handle_retry with return, no yield
-def handle_retry(history, retry_data, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-    update_seed()
-    new_history = history[:retry_data.index]
-    previous_prompt = history[retry_data.index]["content"]
-    result = agent.run_gradio_chat(
-        new_history + [{"role": "user", "content": previous_prompt}],
-        temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round
-    )
-    # If it's a generator, convert to list to avoid Gradio errors
-    if hasattr(result, "__iter__") and not isinstance(result, (str, dict, list)):
-        result = list(result)
-    return result
-
-# ===== Build Gradio Interface =====
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.Markdown(INTRO)
-
-    temperature = gr.Slider(0, 1, step=0.1, value=0.3, label="Temperature")
-    max_new_tokens = gr.Slider(128, 4096, step=1, value=1024, label="Max New Tokens")
-    max_tokens = gr.Slider(128, 32000, step=1, value=8192, label="Max Total Tokens")
-    max_round = gr.Slider(1, 50, step=1, value=30, label="Max Rounds")
-    multi_agent = gr.Checkbox(label="Enable Multi-agent Reasoning", value=False)
-    conversation_state = gr.State([])
-
-    chatbot = gr.Chatbot(
-        label="TxAgent",
-        placeholder=PLACEHOLDER,
-        height=700,
-        type="messages",
-        show_copy_button=True
-    )
-
-    # ✅ Retry logic added safely
-    chatbot.retry(
-        handle_retry,
-        chatbot, chatbot,
-        temperature, max_new_tokens, max_tokens,
-        multi_agent, conversation_state, max_round
-    )
-
-    gr.ChatInterface(
-        fn=handle_chat,
-        chatbot=chatbot,
-        additional_inputs=[
+# Initialize agent placeholder
+agent = None
+
+# ===== Build Gradio UI =====
+def create_ui():
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(DESCRIPTION)
+        gr.Markdown(INTRO)
+
+        temperature = gr.Slider(0, 1, step=0.1, value=0.3, label="Temperature")
+        max_new_tokens = gr.Slider(128, 4096, step=1, value=1024, label="Max New Tokens")
+        max_tokens = gr.Slider(128, 32000, step=1, value=8192, label="Max Total Tokens")
+        max_round = gr.Slider(1, 50, step=1, value=30, label="Max Rounds")
+        multi_agent = gr.Checkbox(label="Enable Multi-agent Reasoning", value=False)
+        conversation_state = gr.State([])
+
+        chatbot = gr.Chatbot(
+            label="TxAgent",
+            placeholder=PLACEHOLDER,
+            height=700,
+            type="messages",
+            show_copy_button=True
+        )
+
+        # Retry logic
+        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]
+            prompt = history[retry_data.index]["content"]
+            result = agent.run_gradio_chat(
+                new_history + [{"role": "user", "content": prompt}],
+                temperature, max_new_tokens, max_tokens,
+                multi_agent, conversation, max_round
+            )
+            if hasattr(result, "__iter__") and not isinstance(result, (str, dict, list)):
+                result = list(result)
+            return result
+
+        chatbot.retry(
+            handle_retry,
+            chatbot, chatbot,
             temperature, max_new_tokens, max_tokens,
             multi_agent, conversation_state, max_round
-        ],
-        examples=question_examples,
-        css=chat_css,
-        cache_examples=False,
-        fill_height=True,
-        fill_width=True,
-        stop_btn=True
+        )
+
+        # Main 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)
+
+        gr.ChatInterface(
+            fn=handle_chat,
+            chatbot=chatbot,
+            additional_inputs=[
+                temperature, max_new_tokens, max_tokens,
+                multi_agent, conversation_state, max_round
+            ],
+            examples=question_examples,
+            css=chat_css,
+            cache_examples=False,
+            fill_height=True,
+            fill_width=True,
+            stop_btn=True
+        )
+
+        gr.Markdown(LICENSE)
+    return demo
+
+# === VLLM-safe entry point ===
+if __name__ == "__main__":
+    agent = TxAgent(
+        model_name,
+        rag_model_name,
+        tool_files_dict=new_tool_files,
+        force_finish=True,
+        enable_checker=True,
+        step_rag_num=10,
+        seed=100,
+        additional_default_tools=["DirectResponse", "RequireClarification"]
     )
+    agent.init_model()
 
-    gr.Markdown(LICENSE)
-
-# ✅ Ensure launch works on Hugging Face Spaces
-demo.launch()
+    demo = create_ui()
+    demo.launch()