Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import torch
 import json
 import logging
@@ -8,14 +7,14 @@ from importlib.resources import files
 from txagent import TxAgent
 from tooluniverse import ToolUniverse
 
-# Logging setup
+# Setup logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-# Paths and environment
+# Env vars
 current_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -69,19 +68,14 @@ def patch_embedding_loading():
                     logger.error("No method found to access tools from ToolUniverse")
                     return False
 
-                current_count = len(tools)
-                embedding_count = len(self.tool_desc_embedding)
-
-                if current_count != embedding_count:
-                    logger.warning(f"Tool count mismatch (tools: {current_count}, embeddings: {embedding_count})")
-                    if current_count < embedding_count:
-                        self.tool_desc_embedding = self.tool_desc_embedding[:current_count]
-                        logger.info(f"Truncated embeddings to match {current_count} tools")
+                if len(tools) != len(self.tool_desc_embedding):
+                    logger.warning("Tool count and embedding count mismatch.")
+                    if len(tools) < len(self.tool_desc_embedding):
+                        self.tool_desc_embedding = self.tool_desc_embedding[:len(tools)]
                     else:
-                        last_embedding = self.tool_desc_embedding[-1]
-                        padding = [last_embedding] * (current_count - embedding_count)
+                        last_emb = self.tool_desc_embedding[-1]
+                        padding = [last_emb] * (len(tools) - len(self.tool_desc_embedding))
                         self.tool_desc_embedding = torch.cat([self.tool_desc_embedding] + padding)
-                        logger.info(f"Padded embeddings to match {current_count} tools")
 
                 return True
 
@@ -90,26 +84,17 @@ def patch_embedding_loading():
                 return False
 
         ToolRAGModel.load_tool_desc_embedding = patched_load
-        logger.info("Successfully patched embedding loading")
+        logger.info("Successfully patched ToolRAGModel")
 
     except Exception as e:
-        logger.error(f"Failed to patch embedding loading: {str(e)}")
-        raise
+        logger.error(f"Failed to patch embedding loader: {str(e)}")
 
 def prepare_tool_files():
     os.makedirs(os.path.join(current_dir, 'data'), exist_ok=True)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
-        logger.info("Generating tool list using ToolUniverse...")
         try:
             tu = ToolUniverse()
-            if hasattr(tu, 'get_all_tools'):
-                tools = tu.get_all_tools()
-            elif hasattr(tu, 'tools'):
-                tools = tu.tools
-            else:
-                tools = []
-                logger.error("Could not access tools from ToolUniverse")
-
+            tools = tu.get_all_tools() if hasattr(tu, 'get_all_tools') else getattr(tu, 'tools', [])
             with open(CONFIG["tool_files"]["new_tool"], "w") as f:
                 json.dump(tools, f, indent=2)
             logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
@@ -119,7 +104,6 @@ def prepare_tool_files():
 def create_agent():
     patch_embedding_loading()
     prepare_tool_files()
-
     try:
         agent = TxAgent(
             CONFIG["model_name"],
@@ -134,27 +118,23 @@ def create_agent():
         agent.init_model()
         return agent
     except Exception as e:
-        logger.error(f"Failed to create agent: {str(e)}")
+        logger.error(f"Failed to create TxAgent: {str(e)}")
         raise
 
-def respond(message_input, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-    # Ensure the message is a proper string
-    message = message_input.strip() if isinstance(message_input, str) else ""
-    
-    if len(message) <= 10:
-        return chat_history + [["assistant", "Hi, I am TxAgent, an assistant for answering biomedical questions. Please provide a valid message with a string longer than 10 characters."]]
+# ✅ GRADIO 5.x-compatible message format
+def respond(msg, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+    if not isinstance(msg, str) or len(msg.strip()) <= 10:
+        return chat_history + [{"role": "assistant", "content": "Hi, I am TxAgent. Please provide a valid question with more than 10 characters."}]
 
-    # Add user message to history
-    chat_history = chat_history + [["user", message]]
+    chat_history = chat_history + [{"role": "user", "content": msg.strip()}]
 
     print("\n==== DEBUG ====")
-    print("User Message:", message)
+    print("User Message:", msg)
     print("Chat History:", chat_history)
     print("================\n")
 
     try:
-        # Format for model: list of (role, message) tuples
-        formatted_history = [(role, content) for role, content in chat_history]
+        formatted_history = [(m["role"], m["content"]) for m in chat_history]
 
         response_generator = agent.run_gradio_chat(
             formatted_history,
@@ -173,15 +153,15 @@ def respond(message_input, chat_history, temperature, max_new_tokens, max_tokens
             else:
                 collected += str(chunk)
 
-        chat_history.append(["assistant", collected])
+        chat_history.append({"role": "assistant", "content": collected})
     except Exception as e:
-        chat_history.append(["assistant", f"Error: {str(e)}"])
+        chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
 
     return chat_history
 
 def create_demo(agent):
     with gr.Blocks(css=chat_css) as demo:
-        chatbot = gr.Chatbot(label="TxAgent", render_markdown=True)
+        chatbot = gr.Chatbot(label="TxAgent", type="messages", render_markdown=True)
         msg = gr.Textbox(label="Your question", placeholder="Type your biomedical query...", scale=6)
         with gr.Row():
             temp = gr.Slider(0, 1, value=0.3, label="Temperature")
@@ -205,7 +185,7 @@ def main():
         global agent
         agent = create_agent()
         demo = create_demo(agent)
-        demo.launch()
+        demo.launch(share=False)  # Set to True to get a public link
     except Exception as e:
         logger.error(f"Application failed to start: {str(e)}")
         raise