Update app.py

app.py

@@ -33,45 +33,54 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
-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)
+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__
+        
+        def patched_init(self, *args, **kwargs):
+            # First let the original initialization happen
+            original_init(self, *args, **kwargs)
             
-            # 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")
+            # Then handle the embeddings our way
+            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:
-                    # 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")
+                    logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
             
-            return True
-            
-        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
+            except Exception as e:
+                logger.error(f"Failed to load embeddings: {str(e)}")
+        
+        # Apply the patch
+        TxAgentClass.__init__ = patched_init
+        logger.info("Successfully patched embedding loading")
+        
+    except Exception as e:
+        logger.error(f"Failed to patch embedding loading: {str(e)}")
+        raise
 
 class TxAgentApp:
     def __init__(self):
@@ -84,7 +93,7 @@ class TxAgentApp:
         
         try:
             # Apply our patch before initialization
-            patch_toolrag_class()
+            patch_embedding_loading()
             
             logger.info("Initializing TxAgent...")
             self.agent = TxAgent(
@@ -149,7 +158,11 @@ def create_interface():
             init_btn = gr.Button("Initialize Model", variant="primary")
             init_status = gr.Textbox(label="Status", interactive=False)
         
-        chatbot = gr.Chatbot(height=500, label="Conversation")
+        chatbot = gr.Chatbot(
+            height=500, 
+            label="Conversation",
+            type="messages"  # Fixing the deprecation warning
+        )
         msg = gr.Textbox(label="Your clinical question")
         clear_btn = gr.Button("Clear Chat")