Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
-from txagent import TxAgent
-import logging
 import os
+import logging
+from txagent import TxAgent
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -18,11 +18,19 @@ class TxAgentApp:
             agent = TxAgent(
                 model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
                 rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-                # Remove unsupported parameters
-                force_finish=True,
-                enable_checker=True,
+                enable_finish=True,
+                enable_rag=True,
+                enable_summary=False,
+                init_rag_num=0,
                 step_rag_num=10,
+                summary_mode='step',
+                summary_skip_last_k=0,
+                summary_context_length=None,
+                force_finish=True,
+                avoid_repeat=True,
                 seed=42,
+                enable_checker=True,
+                enable_chat=False,
                 additional_default_tools=["DirectResponse", "RequireClarification"]
             )
             logger.info("Model loading complete")
@@ -31,87 +39,117 @@ class TxAgentApp:
             logger.error(f"Initialization failed: {str(e)}")
             raise
 
-    def respond(self, message, history):
-        """Handle streaming responses with Gradio 5.23"""
+    def respond(self, message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round):
+        """Handle streaming responses with Gradio"""
         try:
+            if not isinstance(message, str) or len(message.strip()) <= 10:
+                yield "Please provide a valid message longer than 10 characters."
+                return
+
             response_generator = self.agent.run_gradio_chat(
-                message=message,
+                message=message.strip(),
                 history=history,
-                temperature=0.3,
-                max_new_tokens=2048,
-                stream=True
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                max_token=max_tokens,
+                call_agent=multi_agent,
+                conversation=conversation_state,
+                max_round=max_round,
+                seed=42
             )
-            
+
+            full_response = ""
             for chunk in response_generator:
-                if isinstance(chunk, dict):
-                    yield chunk.get("content", "")
+                if isinstance(chunk, dict) and "content" in chunk:
+                    content = chunk["content"]
                 elif isinstance(chunk, str):
-                    yield chunk
+                    content = chunk
+                else:
+                    content = str(chunk)
+                
+                full_response += content
+                yield full_response
+
         except Exception as e:
-            logger.error(f"Generation error: {str(e)}")
+            logger.error(f"Error in respond function: {str(e)}")
             yield f"⚠️ Error: {str(e)}"
 
-# Initialize the app
-app = TxAgentApp()
-
-# Gradio 5.23 interface
-with gr.Blocks(
-    title="TxAgent Medical AI",
-    theme=gr.themes.Soft(spacing_size="sm", radius_size="none")
-) as demo:
-    gr.Markdown("""<h1 style='text-align: center'>TxAgent Biomedical Assistant</h1>""")
+def create_demo():
+    """Create and configure the Gradio interface"""
+    app = TxAgentApp()
     
-    with gr.Row(equal_height=False):
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                height=650,
-                bubble_full_width=False
-            )
-        with gr.Column(scale=1):
-            with gr.Accordion("⚙️ Parameters", open=False):
-                temperature = gr.Slider(0, 1, value=0.3, label="Creativity")
-                max_tokens = gr.Slider(128, 4096, value=1024, step=128, label="Max Response Length")
-                rag_toggle = gr.Checkbox(value=True, label="Enable RAG")
-            
-            msg = gr.Textbox(
-                label="Your medical query",
-                placeholder="Enter your biomedical question...",
-                lines=5,
-                max_lines=10
-            )
-            submit_btn = gr.Button("Submit", variant="primary")
-            clear_btn = gr.Button("Clear History")
-
-    # Chat interface
-    msg.submit(
-        app.respond,
-        [msg, chatbot],
-        chatbot,
-        api_name="chat"
-    ).then(
-        lambda: "", None, msg
-    )
-
-    submit_btn.click(
-        app.respond,
-        [msg, chatbot],
-        chatbot,
-        api_name="chat"
-    ).then(
-        lambda: "", None, msg
-    )
+    with gr.Blocks(
+        title="TxAgent Medical AI",
+        theme=gr.themes.Soft(spacing_size="sm", radius_size="none")
+    ) as demo:
+        gr.Markdown("""<h1 style='text-align: center'>TxAgent Biomedical Assistant</h1>""")
+        
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot(
+                    height=650,
+                    bubble_full_width=False,
+                    render_markdown=True
+                )
+                msg = gr.Textbox(
+                    label="Your medical query",
+                    placeholder="Enter your biomedical question...",
+                    lines=5,
+                    max_lines=10
+                )
+                
+            with gr.Column(scale=1):
+                with gr.Accordion("⚙️ Parameters", open=False):
+                    temperature = gr.Slider(0, 1, value=0.3, label="Creativity")
+                    max_new_tokens = gr.Slider(128, 4096, value=1024, step=128, label="Max Response Length")
+                    max_tokens = gr.Slider(128, 81920, value=81920, step=1024, label="Max Total Tokens")
+                    max_rounds = gr.Slider(1, 30, value=10, step=1, label="Max Rounds")
+                    multi_agent = gr.Checkbox(value=False, label="Multi-Agent Mode")
+                
+                submit_btn = gr.Button("Submit", variant="primary")
+                clear_btn = gr.Button("Clear History")
+        
+        conversation_state = gr.State([])
+        
+        # Chat interface
+        msg.submit(
+            app.respond,
+            [msg, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+            chatbot
+        ).then(lambda: "", None, msg)
+        
+        submit_btn.click(
+            app.respond,
+            [msg, chatbot, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
+            chatbot
+        ).then(lambda: "", None, msg)
+        
+        clear_btn.click(
+            lambda: ([], []), 
+            None, 
+            [chatbot, conversation_state]
+        )
+    
+    return demo
 
-    clear_btn.click(
-        lambda: [], None, chatbot
-    )
+def main():
+    """Main entry point for the application"""
+    try:
+        logger.info("Starting TxAgent application...")
+        demo = create_demo()
+        
+        logger.info("Launching Gradio interface...")
+        demo.queue(
+            concurrency_count=5,
+            max_size=20
+        ).launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False
+        )
+    except Exception as e:
+        logger.error(f"Application failed to start: {str(e)}")
+        raise
 
-# Launch configuration
 if __name__ == "__main__":
-    demo.queue(
-        concurrency_count=5,
-        max_size=20
-    ).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    main()