Update app.py

app.py

@@ -1,47 +1,54 @@
 import os
+import json
+import logging
 import torch
-import requests
-from huggingface_hub import hf_hub_download, snapshot_download
 from txagent import TxAgent
 import gradio as gr
+from huggingface_hub import hf_hub_download, snapshot_download
+from tooluniverse import ToolUniverse
 
 # Configuration
 CONFIG = {
     "model_name": "mims-harvard/TxAgent-T1-Llama-3.1-8B",
     "rag_model_name": "mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_e27fb393f3144ec28f620f33d4d79911.pt",
+    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding.pt",
     "local_dir": "./models",
     "tool_files": {
-        'new_tool': './data/new_tool.json',
-        'opentarget': './data/opentarget_tools.json',
-        'fda_drug_label': './data/fda_drug_labeling_tools.json',
-        'special_tools': './data/special_tools.json',
-        'monarch': './data/monarch_tools.json'
+        "new_tool": "./data/new_tool.json"
     }
 }
 
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def prepare_tool_files():
+    os.makedirs("./data", exist_ok=True)
+    if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
+        logger.info("Generating tool list using ToolUniverse...")
+        tu = ToolUniverse()
+        tools = tu.get_all_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']}")
+
+
 def download_model_files():
-    """Download all required model files from Hugging Face Hub"""
     os.makedirs(CONFIG["local_dir"], exist_ok=True)
-    os.makedirs("./data", exist_ok=True)
-    
     print("Downloading model files...")
-    
-    # Download main model
+
     snapshot_download(
         repo_id=CONFIG["model_name"],
         local_dir=os.path.join(CONFIG["local_dir"], CONFIG["model_name"]),
         resume_download=True
     )
-    
-    # Download RAG model
+
     snapshot_download(
         repo_id=CONFIG["rag_model_name"],
         local_dir=os.path.join(CONFIG["local_dir"], CONFIG["rag_model_name"]),
         resume_download=True
     )
-    
-    # Try to download the embeddings file
+
     try:
         hf_hub_download(
             repo_id=CONFIG["rag_model_name"],
@@ -55,30 +62,20 @@ def download_model_files():
         print("Will attempt to generate it instead")
 
 def generate_embeddings(agent):
-    """Generate and save tool embeddings if missing"""
     embedding_path = os.path.join(CONFIG["local_dir"], CONFIG["embedding_filename"])
-    
+
     if os.path.exists(embedding_path):
         print("Embeddings file already exists")
         return
-    
+
     print("Generating missing tool embeddings...")
-    
     try:
-        # Get all tools from the tool universe
         tools = agent.tooluniverse.get_all_tools()
-        tool_descriptions = [tool['description'] for tool in tools]
-        
-        # Generate embeddings using the RAG model
-        embeddings = agent.rag_model.generate_embeddings(tool_descriptions)
-        
-        # Save the embeddings
+        descriptions = [tool["description"] for tool in tools]
+        embeddings = agent.rag_model.generate_embeddings(descriptions)
         torch.save(embeddings, embedding_path)
-        print(f"Embeddings saved to {embedding_path}")
-        
-        # Update the RAG model to use the new embeddings
         agent.rag_model.tool_desc_embedding = embeddings
-        
+        print(f"Embeddings saved to {embedding_path}")
     except Exception as e:
         print(f"Failed to generate embeddings: {e}")
         raise
@@ -91,9 +88,8 @@ class TxAgentApp:
     def initialize(self):
         if self.is_initialized:
             return "Already initialized"
-        
+
         try:
-            # Initialize the agent
             self.agent = TxAgent(
                 CONFIG["model_name"],
                 CONFIG["rag_model_name"],
@@ -102,37 +98,24 @@ class TxAgentApp:
                 enable_checker=True,
                 step_rag_num=10,
                 seed=100,
-                additional_default_tools=['DirectResponse', 'RequireClarification']
+                additional_default_tools=["DirectResponse", "RequireClarification"]
             )
-            
-            # Initialize model
             self.agent.init_model()
-            
-            # Handle embeddings
             generate_embeddings(self.agent)
-            
             self.is_initialized = True
-            return "TxAgent initialized successfully"
-            
+            return "✅ TxAgent initialized successfully"
         except Exception as e:
-            return f"Initialization failed: {str(e)}"
+            return f"❌ Initialization failed: {str(e)}"
 
     def chat(self, message, history):
         if not self.is_initialized:
-            return history + [(message, "Error: Please initialize the model first")]
-        
+            return history + [(message, "⚠️ Error: Model not initialized")]
+
         try:
-            # Convert history to messages format
-            messages = []
-            for user_msg, bot_msg in history:
-                messages.append({"role": "user", "content": user_msg})
-                messages.append({"role": "assistant", "content": bot_msg})
-            messages.append({"role": "user", "content": message})
-            
-            # Get response
             response = ""
             for chunk in self.agent.run_gradio_chat(
-                messages,
+                message=message,
+                history=history,
                 temperature=0.3,
                 max_new_tokens=1024,
                 max_tokens=8192,
@@ -141,28 +124,24 @@ class TxAgentApp:
                 max_round=30
             ):
                 response += chunk
-            
+
             return history + [(message, response)]
         except Exception as e:
             return history + [(message, f"Error: {str(e)}")]
 
 def create_interface():
     app = TxAgentApp()
-    
     with gr.Blocks(title="TxAgent") as demo:
-        gr.Markdown("# TxAgent: Therapeutic Reasoning AI")
-        
-        # Initialization
+        gr.Markdown("# 🧠 TxAgent: Therapeutic Reasoning AI")
+
         with gr.Row():
             init_btn = gr.Button("Initialize Model", variant="primary")
             init_status = gr.Textbox(label="Initialization Status")
-        
-        # Chat interface
-        chatbot = gr.Chatbot(height=600)
+
+        chatbot = gr.Chatbot(height=600, label="Conversation")
         msg = gr.Textbox(label="Your Question")
         submit_btn = gr.Button("Submit")
-        
-        # Examples
+
         gr.Examples(
             examples=[
                 "How to adjust Journavx dosage for hepatic impairment?",
@@ -171,29 +150,15 @@ def create_interface():
             ],
             inputs=msg
         )
-        
-        # Event handlers
-        init_btn.click(
-            app.initialize,
-            outputs=init_status
-        )
-        
-        def respond(message, chat_history):
-            return app.chat(message, chat_history)
-        
-        msg.submit(respond, [msg, chatbot], chatbot)
-        submit_btn.click(respond, [msg, chatbot], chatbot)
-    
+
+        init_btn.click(fn=app.initialize, outputs=init_status)
+        msg.submit(fn=app.chat, inputs=[msg, chatbot], outputs=chatbot)
+        submit_btn.click(fn=app.chat, inputs=[msg, chatbot], outputs=chatbot)
+
     return demo
 
 if __name__ == "__main__":
-    # First download all required files
+    prepare_tool_files()
     download_model_files()
-    
-    # Then create and launch the interface
     interface = create_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True
-    )
+    interface.launch(server_name="0.0.0.0", server_port=7860, share=False)