Update app.py

app.py

@@ -6,11 +6,11 @@ from txagent import TxAgent
 import gradio as gr
 from tooluniverse import ToolUniverse
 
-# Configuration - Using your existing embedding file
+# Configuration with hardcoded embedding 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_e27fb393f3144ec28f620f33d4d79911.pt",
+    "embedding_filename": "ToolRAG-T1-GTE-Qwen2-1.5Btool_embedding_47dc56b3e3ddeb31af4f19defdd538d984de1500368852a0fab80bc2e826c944.pt",
     "tool_files": {
         "new_tool": "./data/new_tool.json"
     }
@@ -33,21 +33,45 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
-def load_embeddings(agent):
-    embedding_path = CONFIG["embedding_filename"]
-    try:
-        if os.path.exists(embedding_path):
-            logger.info(f"✅ Loading existing embeddings from {embedding_path}")
-            embeddings = torch.load(embedding_path)
-            agent.rag_model.tool_desc_embedding = embeddings
+def patch_toolrag_class():
+    """Monkey-patch the ToolRAG class to use our embedding file and handle tool count mismatch"""
+    from txagent.toolrag import ToolRAG
+    
+    original_load = ToolRAG.load_tool_desc_embedding
+    
+    def patched_load(self, tooluniverse):
+        try:
+            # Load our specific embedding file
+            self.tool_desc_embedding = torch.load(CONFIG["embedding_filename"])
+            
+            # Get current tools and their count
+            tools = tooluniverse.get_all_tools()
+            current_tool_count = len(tools)
+            embedding_count = len(self.tool_desc_embedding)
+            
+            # If counts don't match, truncate or pad as needed
+            if current_tool_count != embedding_count:
+                logger.warning(f"Tool count mismatch! Tools: {current_tool_count}, Embeddings: {embedding_count}")
+                
+                if current_tool_count < embedding_count:
+                    # Truncate embeddings to match tool count
+                    self.tool_desc_embedding = self.tool_desc_embedding[:current_tool_count]
+                    logger.warning(f"Truncated embeddings to {current_tool_count} vectors")
+                else:
+                    # Pad with zeros (last embedding) if tools > embeddings
+                    last_embedding = self.tool_desc_embedding[-1]
+                    padding = [last_embedding] * (current_tool_count - embedding_count)
+                    self.tool_desc_embedding = torch.cat([self.tool_desc_embedding] + padding)
+                    logger.warning(f"Padded embeddings with {current_tool_count - embedding_count} vectors")
+            
             return True
-        else:
-            logger.error(f"❌ Embedding file not found at {embedding_path}")
-            logger.info("Please ensure the embedding file is in the root directory")
+            
+        except Exception as e:
+            logger.error(f"Failed to load embeddings: {str(e)}")
             return False
-    except Exception as e:
-        logger.error(f"Failed to load embeddings: {str(e)}")
-        return False
+    
+    # Apply the patch
+    ToolRAG.load_tool_desc_embedding = patched_load
 
 class TxAgentApp:
     def __init__(self):
@@ -59,9 +83,10 @@ class TxAgentApp:
             return "✅ Already initialized"
         
         try:
-            logger.info("Initializing TxAgent...")
+            # Apply our patch before initialization
+            patch_toolrag_class()
             
-            # Initialize TxAgent
+            logger.info("Initializing TxAgent...")
             self.agent = TxAgent(
                 CONFIG["model_name"],
                 CONFIG["rag_model_name"],
@@ -73,15 +98,9 @@ class TxAgentApp:
                 additional_default_tools=["DirectResponse", "RequireClarification"]
             )
             
-            # Initialize models
             logger.info("Loading models...")
             self.agent.init_model()
             
-            # Load embeddings
-            logger.info("Loading embeddings...")
-            if not load_embeddings(self.agent):
-                return "❌ Failed to load embeddings - check logs"
-            
             self.is_initialized = True
             return "✅ TxAgent initialized successfully"
             
@@ -123,6 +142,7 @@ def create_interface():
     ) as demo:
         gr.Markdown("""
         # 🧠 TxAgent: Therapeutic Reasoning AI
+        ### (Using pre-loaded embeddings)
         """)
         
         with gr.Row():