Update app.py

app.py

@@ -1,23 +1,21 @@
 import os
-import torch
 import json
+import torch
 import logging
 import gradio as gr
 from importlib.resources import files
 from txagent import TxAgent
 from tooluniverse import ToolUniverse
 
-# Setup logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-# Env vars
-current_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
+current_dir = os.path.dirname(os.path.abspath(__file__))
 
 CONFIG = {
     "model_name": "mims-harvard/TxAgent-T1-Llama-3.1-8B",
@@ -32,74 +30,46 @@ CONFIG = {
     }
 }
 
-chat_css = """
-.gr-button { font-size: 20px !important; }
-.gr-button svg { width: 32px !important; height: 32px !important; }
-"""
-
-def safe_load_embeddings(filepath: str) -> any:
+def safe_load_embeddings(filepath):
     try:
         return torch.load(filepath, weights_only=True)
     except Exception as e:
-        logger.warning(f"Secure load failed, trying with weights_only=False: {str(e)}")
+        logger.warning(f"Retrying with weights_only=False due to: {e}")
         try:
             return torch.load(filepath, weights_only=False)
         except Exception as e:
-            logger.error(f"Failed to load embeddings: {str(e)}")
+            logger.error(f"Failed to load embeddings: {e}")
             return None
 
 def patch_embedding_loading():
-    try:
-        from txagent.toolrag import ToolRAGModel
-
-        def patched_load(self, tooluniverse):
-            try:
-                if not os.path.exists(CONFIG["embedding_filename"]):
-                    logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
-                    return False
-
-                self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
-
-                if hasattr(tooluniverse, 'get_all_tools'):
-                    tools = tooluniverse.get_all_tools()
-                elif hasattr(tooluniverse, 'tools'):
-                    tools = tooluniverse.tools
-                else:
-                    logger.error("No method found to access tools from ToolUniverse")
-                    return False
-
-                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_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)
-
-                return True
-
-            except Exception as e:
-                logger.error(f"Failed to load embeddings: {str(e)}")
+    from txagent.toolrag import ToolRAGModel
+    def patched_load(self, tooluniverse):
+        try:
+            if not os.path.exists(CONFIG["embedding_filename"]):
                 return False
+            self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
 
-        ToolRAGModel.load_tool_desc_embedding = patched_load
-        logger.info("Successfully patched ToolRAGModel")
+            tools = tooluniverse.get_all_tools() if hasattr(tooluniverse, "get_all_tools") else getattr(tooluniverse, "tools", [])
+            if len(tools) != len(self.tool_desc_embedding):
+                logger.warning("Tool count mismatch.")
+                self.tool_desc_embedding = self.tool_desc_embedding[:len(tools)]
+            return True
+        except Exception as e:
+            logger.error(f"Embedding load failed: {e}")
+            return False
 
-    except Exception as e:
-        logger.error(f"Failed to patch embedding loader: {str(e)}")
+    ToolRAGModel.load_tool_desc_embedding = patched_load
 
 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"]):
         try:
             tu = ToolUniverse()
-            tools = tu.get_all_tools() if hasattr(tu, 'get_all_tools') else getattr(tu, 'tools', [])
+            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']}")
         except Exception as e:
-            logger.error(f"Failed to prepare tool files: {str(e)}")
+            logger.error(f"Tool generation failed: {e}")
 
 def create_agent():
     patch_embedding_loading()
@@ -112,83 +82,69 @@ def create_agent():
             force_finish=True,
             enable_checker=True,
             step_rag_num=10,
-            seed=100,
-            additional_default_tools=['DirectResponse', 'RequireClarification']
+            seed=42,
+            additional_default_tools=["DirectResponse", "RequireClarification"]
         )
         agent.init_model()
         return agent
     except Exception as e:
-        logger.error(f"Failed to create TxAgent: {str(e)}")
+        logger.error(f"Agent initialization failed: {e}")
         raise
 
-# ✅ GRADIO 5.x-compatible message format
+# ✅ FIXED: Proper message formatting
 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."}]
+        return chat_history + [{"role": "assistant", "content": "Hi, I am TxAgent. Please provide a valid message longer than 10 characters."}]
 
-    chat_history = chat_history + [{"role": "user", "content": msg.strip()}]
-
-    print("\n==== DEBUG ====")
-    print("User Message:", msg)
-    print("Chat History:", chat_history)
-    print("================\n")
+    message = msg.strip()
+    chat_history.append({"role": "user", "content": message})
+    formatted_history = [(m["role"], m["content"]) for m in chat_history]
 
     try:
-        formatted_history = [(m["role"], m["content"]) for m in chat_history]
-
         response_generator = agent.run_gradio_chat(
-            formatted_history,
-            temperature,
-            max_new_tokens,
-            max_tokens,
-            multi_agent,
-            conversation,
-            max_round
+            message=message,
+            history=formatted_history,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            max_token=max_tokens,
+            call_agent=multi_agent,
+            conversation=conversation,
+            max_round=max_round,
+            seed=42,
+            call_agent_level=None,
+            sub_agent_task=None
         )
-
         collected = ""
         for chunk in response_generator:
-            if isinstance(chunk, dict):
-                collected += chunk.get("content", "")
-            else:
-                collected += str(chunk)
-
+            collected += chunk.get("content", "") if isinstance(chunk, dict) else str(chunk)
         chat_history.append({"role": "assistant", "content": collected})
     except Exception as e:
-        chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
-
+        chat_history.append({"role": "assistant", "content": f"Error: {e}"})
     return chat_history
 
 def create_demo(agent):
-    with gr.Blocks(css=chat_css) as demo:
+    with gr.Blocks(css=".gr-button { font-size: 18px !important; }") as demo:
         chatbot = gr.Chatbot(label="TxAgent", type="messages", render_markdown=True)
-        msg = gr.Textbox(label="Your question", placeholder="Type your biomedical query...", scale=6)
+        msg = gr.Textbox(label="Your question", placeholder="Ask a biomedical question...", scale=6)
         with gr.Row():
             temp = gr.Slider(0, 1, value=0.3, label="Temperature")
             max_new_tokens = gr.Slider(128, 4096, value=1024, label="Max New Tokens")
             max_tokens = gr.Slider(128, 81920, value=81920, label="Max Total Tokens")
             max_rounds = gr.Slider(1, 30, value=30, label="Max Rounds")
             multi_agent = gr.Checkbox(label="Multi-Agent Mode")
-        with gr.Row():
-            submit = gr.Button("Ask TxAgent")
-
+        submit = gr.Button("Ask TxAgent")
         submit.click(
             respond,
             inputs=[msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, gr.State([]), max_rounds],
             outputs=[chatbot]
         )
-
     return demo
 
 def main():
-    try:
-        global agent
-        agent = create_agent()
-        demo = create_demo(agent)
-        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
+    global agent
+    agent = create_agent()
+    demo = create_demo(agent)
+    demo.launch(share=False)
 
 if __name__ == "__main__":
     main()