Update app.py

app.py

@@ -64,7 +64,7 @@ def patch_embedding_loading():
                 self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
                 
                 # Handle tool count mismatch
-                tools = tooluniverse.get_all_tools()
+                tools = tooluniverse.tools  # Changed from get_all_tools()
                 current_count = len(tools)
                 embedding_count = len(self.tool_desc_embedding)
                 
@@ -133,26 +133,32 @@ class TxAgentApp:
 
     def chat(self, message, history):
         if not self.is_initialized:
-            return history + [(message, "⚠️ Please initialize the model first")]
+            return {"role": "assistant", "content": "⚠️ Please initialize the model first"}
         
         try:
             response = ""
+            # Modified to use the correct parameter name (max_length instead of max_tokens)
             for chunk in self.agent.run_gradio_chat(
                 message=message,
                 history=history,
                 temperature=0.3,
                 max_new_tokens=1024,
-                max_tokens=8192,
+                max_length=8192,  # Changed from max_tokens
                 multi_agent=False,
                 conversation=[],
                 max_round=30
             ):
                 response += chunk
-                yield history + [(message, response)]
+            
+            # Format response in the expected messages format
+            return [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": response}
+            ]
                 
         except Exception as e:
             logger.error(f"Chat error: {str(e)}")
-            yield history + [(message, f"Error: {str(e)}")]
+            return {"role": "assistant", "content": f"Error: {str(e)}"}
 
 def create_interface():
     app = TxAgentApp()
@@ -175,7 +181,7 @@ def create_interface():
         chatbot = gr.Chatbot(
             height=500, 
             label="Conversation",
-            type="messages"  # Fixing the deprecation warning
+            bubble_full_width=False
         )
         msg = gr.Textbox(label="Your clinical question")
         clear_btn = gr.Button("Clear Chat")
@@ -189,13 +195,23 @@ def create_interface():
             inputs=msg
         )
         
+        def wrapper_initialize():
+            status = app.initialize()
+            return status, gr.update(interactive=False)
+        
+        def wrapper_chat(message, chat_history):
+            response = app.chat(message, chat_history)
+            if isinstance(response, dict):  # Error case
+                return chat_history + [response]
+            return response  # Normal case
+        
         init_btn.click(
-            fn=app.initialize,
-            outputs=init_status
+            fn=wrapper_initialize,
+            outputs=[init_status, init_btn]
         )
         
         msg.submit(
-            fn=app.chat,
+            fn=wrapper_chat,
             inputs=[msg, chatbot],
             outputs=chatbot
         )