Update app.py

app.py

@@ -115,24 +115,28 @@ def patch_embedding_loading():
     """Monkey-patch the embedding loading functionality"""
     try:
         from txagent.toolrag import ToolRAGModel
-        
+
         original_load = ToolRAGModel.load_tool_desc_embedding
-        
+
         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"])
-                
-                tools = tooluniverse.get_all_tools()
+                if self.tool_desc_embedding is None:
+                    logger.error("Embedding is None, aborting.")
+                    return False
+
+                # Ensure tools is a list (in case it's a generator)
+                tools = list(tooluniverse.get_all_tools()) if hasattr(tooluniverse, 'get_all_tools') else []
                 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")
@@ -141,51 +145,20 @@ def patch_embedding_loading():
                         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
-        
+
         ToolRAGModel.load_tool_desc_embedding = patched_load
         logger.info("Successfully patched embedding loading")
-        
+
     except Exception as e:
         logger.error(f"Failed to patch embedding loading: {str(e)}")
         raise
 
-def prepare_tool_files():
-    """Ensure tool files exist and are populated"""
-    os.makedirs(os.path.join(current_dir, '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() if hasattr(tu, 'get_all_tools') else []
-        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 create_agent():
-    """Create and initialize the TxAgent"""
-    # Apply the embedding patch before creating the agent
-    patch_embedding_loading()
-    prepare_tool_files()
-
-    # Initialize the agent
-    agent = TxAgent(
-        CONFIG["model_name"],
-        CONFIG["rag_model_name"],
-        tool_files_dict=CONFIG["tool_files"],
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=10,
-        seed=100,
-        additional_default_tools=['DirectResponse', 'RequireClarification']
-    )
-    agent.init_model()
-    return agent
-
 def update_model_parameters(agent, enable_finish, enable_rag, enable_summary,
                           init_rag_num, step_rag_num, skip_last_k,
                           summary_mode, summary_skip_last_k, summary_context_length,