Update app.py

app.py

@@ -30,7 +30,6 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
 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.pt",
     "tool_files": {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
         "fda_drug_label": str(files('tooluniverse.data').joinpath('fda_drug_labeling_tools.json')),
@@ -40,16 +39,6 @@ CONFIG = {
     }
 }
 
-def generate_tool_embeddings(agent):
-    tu = ToolUniverse(tool_files=CONFIG["tool_files"])
-    tu.load_tools()
-    embedding_tensor = agent.rag_model.load_tool_desc_embedding(tu)
-    if embedding_tensor is not None:
-        torch.save(embedding_tensor, CONFIG["embedding_filename"])
-        logger.info(f"Saved new embedding tensor to {CONFIG['embedding_filename']}")
-    else:
-        logger.warning("Embedding generation returned None")
-
 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"]):
@@ -74,8 +63,6 @@ def create_agent():
             seed=42,
             additional_default_tools=["DirectResponse", "RequireClarification"]
         )
-        if not os.path.exists(CONFIG["embedding_filename"]):
-            generate_tool_embeddings(agent)
         agent.init_model()
         return agent
     except Exception as e:
@@ -88,7 +75,7 @@ def respond(msg, chat_history, temperature, max_new_tokens, max_tokens, multi_ag
 
     message = msg.strip()
     chat_history.append({"role": "user", "content": message})
-    formatted_history = chat_history  # format as list of dicts for run_gradio_chat
+    formatted_history = [(m["role"], m["content"]) for m in chat_history if "role" in m and "content" in m]
 
     try:
         response_generator = agent.run_gradio_chat(
@@ -101,20 +88,18 @@ def respond(msg, chat_history, temperature, max_new_tokens, max_tokens, multi_ag
             conversation=conversation,
             max_round=max_round,
             seed=42,
-            call_agent_level=0,
+            call_agent_level=None,
             sub_agent_task=None
         )
 
         collected = ""
         for chunk in response_generator:
-            if isinstance(chunk, list):
-                for msg in chunk:
-                    if isinstance(msg, dict) and "content" in msg:
-                        collected += msg["content"]
-            elif isinstance(chunk, dict) and "content" in chunk:
+            if isinstance(chunk, dict) and "content" in chunk:
                 collected += chunk["content"]
             elif isinstance(chunk, str):
                 collected += chunk
+            elif chunk is not None:
+                collected += str(chunk)
 
         chat_history.append({"role": "assistant", "content": collected or "⚠️ No content returned."})
 
@@ -145,7 +130,8 @@ def main():
     global agent
     agent = create_agent()
     demo = create_demo(agent)
-    demo.launch(share=True)
+    print("Exiting after embedding generation. Please restart the Space manually.")
+    exit()
 
 if __name__ == "__main__":
     main()