Update app.py

app.py

@@ -5,6 +5,10 @@ import torch
 from txagent import TxAgent
 import gradio as gr
 from tooluniverse import ToolUniverse
+import warnings
+
+# Suppress specific warnings
+warnings.filterwarnings("ignore", category=UserWarning)
 
 # Configuration with hardcoded embedding file
 CONFIG = {
@@ -33,49 +37,59 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
+def safe_load_embeddings(filepath):
+    """Safely load embeddings with proper weights_only handling"""
+    try:
+        # First try with weights_only=True (secure mode)
+        return torch.load(filepath, weights_only=True)
+    except Exception as e:
+        logger.warning(f"Secure load failed, trying with weights_only=False: {str(e)}")
+        # If that fails, try with weights_only=False (less secure)
+        return torch.load(filepath, weights_only=False)
+
 def patch_embedding_loading():
     """Monkey-patch the embedding loading functionality"""
     try:
-        # Try to get the RAG model class dynamically
-        from txagent.txagent import TxAgent as TxAgentClass
-        original_init = TxAgentClass.__init__
+        from txagent.toolrag import ToolRAGModel
         
-        def patched_init(self, *args, **kwargs):
-            # First let the original initialization happen
-            original_init(self, *args, **kwargs)
-            
-            # Then handle the embeddings our way
+        original_load = ToolRAGModel.load_tool_desc_embedding
+        
+        def patched_load(self, tooluniverse):
             try:
-                if os.path.exists(CONFIG["embedding_filename"]):
-                    logger.info(f"Loading embeddings from {CONFIG['embedding_filename']}")
-                    self.rag_model.tool_desc_embedding = torch.load(CONFIG["embedding_filename"])
-                    
-                    # Handle tool count mismatch
-                    tools = self.tooluniverse.get_all_tools()
-                    current_count = len(tools)
-                    embedding_count = len(self.rag_model.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.rag_model.tool_desc_embedding = self.rag_model.tool_desc_embedding[:current_count]
-                            logger.info(f"Truncated embeddings to match {current_count} tools")
-                        else:
-                            last_embedding = self.rag_model.tool_desc_embedding[-1]
-                            padding = [last_embedding] * (current_count - embedding_count)
-                            self.rag_model.tool_desc_embedding = torch.cat(
-                                [self.rag_model.tool_desc_embedding] + padding
-                            )
-                            logger.info(f"Padded embeddings to match {current_count} tools")
-                else:
+                if not os.path.exists(CONFIG["embedding_filename"]):
                     logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
-            
+                    return False
+                
+                # Load embeddings safely
+                self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
+                
+                # Handle tool count mismatch
+                tools = tooluniverse.get_all_tools()
+                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")
+                    else:
+                        last_embedding = self.tool_desc_embedding[-1]
+                        padding = [last_embedding] * (current_count - embedding_count)
+                        self.tool_desc_embedding = torch.cat(
+                            [self.tool_desc_embedding] + padding
+                        )
+                        logger.info(f"Padded embeddings to match {current_count} tools")
+                
+                return True
+                
             except Exception as e:
                 logger.error(f"Failed to load embeddings: {str(e)}")
+                return False
         
         # Apply the patch
-        TxAgentClass.__init__ = patched_init
+        ToolRAGModel.load_tool_desc_embedding = patched_load
         logger.info("Successfully patched embedding loading")
         
     except Exception as e: