Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import gradio as gr
 from txagent import TxAgent
+import torch
 
 # ========== Configuration ==========
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -17,7 +18,8 @@ MODEL_CONFIG = {
         'force_finish': True,
         'enable_checker': True,
         'step_rag_num': 10,
-        'seed': 100
+        'seed': 100,
+        'enable_rag': False  # Disable RAG until we resolve the embedding issue
     }
 }
 
@@ -55,7 +57,7 @@ class TxAgentApplication:
             return "Model already initialized"
         
         try:
-            # Initialize the agent
+            # Initialize the agent with RAG disabled initially
             self.agent = TxAgent(
                 MODEL_CONFIG['model_name'],
                 MODEL_CONFIG['rag_model_name'],
@@ -67,11 +69,18 @@ class TxAgentApplication:
             try:
                 self.agent.init_model()
             except Exception as e:
-                # Handle specific tool embedding error
                 if "No such file or directory" in str(e) and "tool_embedding" in str(e):
-                    return ("Error: Missing tool embedding file. "
-                           "Please ensure the RAG model files are properly downloaded.")
-                raise
+                    # Try to generate embeddings if missing
+                    try:
+                        self._generate_missing_embeddings()
+                        # Retry initialization
+                        self.agent.init_model()
+                    except Exception as gen_e:
+                        return (f"Error: Could not generate missing embeddings. "
+                               f"Please ensure all model files are properly downloaded. "
+                               f"Technical details: {str(gen_e)}")
+                else:
+                    raise
             
             self.is_initialized = True
             self.initialization_error = None
@@ -81,6 +90,25 @@ class TxAgentApplication:
             self.initialization_error = str(e)
             return f"Initialization failed: {str(e)}"
 
+    def _generate_missing_embeddings(self):
+        """Attempt to generate missing tool embeddings"""
+        if not hasattr(self.agent, 'rag_model'):
+            raise ValueError("RAG model not initialized")
+        
+        # Get the tools from the tool universe
+        tools = self.agent.tooluniverse.get_all_tools()
+        tool_descriptions = [tool['description'] for tool in tools]
+        
+        # Generate embeddings using the RAG model
+        embeddings = self.agent.rag_model.generate_embeddings(tool_descriptions)
+        
+        # Save the embeddings
+        embedding_path = f"{MODEL_CONFIG['rag_model_name']}_tool_embedding_e27fb393f3144ec28f620f33d4d79911.pt"
+        torch.save(embeddings, embedding_path)
+        
+        # Update the RAG model to use the new embeddings
+        self.agent.rag_model.tool_desc_embedding = embeddings
+
     def chat(self, message, chat_history):
         if not self.is_initialized:
             if self.initialization_error:
@@ -189,5 +217,6 @@ if __name__ == "__main__":
     try:
         interface.queue().launch(**launch_params)
     except Exception as e:
-        print(f"Error launching interface: {e}")
-        interface.launch(**launch_params)  # Fallback without queue
+        print(f"Error launching with queue: {e}")
+        print("Falling back to non-queued launch")
+        interface.launch(**launch_params)