Update memory_logic.py

memory_logic.py

@@ -61,7 +61,7 @@ def _get_sqlite_connection():
     db_dir = os.path.dirname(SQLITE_DB_PATH)
     if db_dir and not os.path.exists(db_dir):
         os.makedirs(db_dir, exist_ok=True)
-    return sqlite3.connect(SQLITE_DB_PATH, timeout=10) # Added timeout
+    return sqlite3.connect(SQLITE_DB_PATH, timeout=10)
 
 def _init_sqlite_tables():
     if STORAGE_BACKEND != "SQLITE" or not sqlite3:
@@ -88,21 +88,48 @@ def _init_sqlite_tables():
     except Exception as e:
         logger.error(f"SQLite table initialization error: {e}", exc_info=True)
 
-# --- Initialization and State Management ---
+
+def _build_faiss_index(items_list, text_extraction_fn):
+    """Builds a FAISS index from a list of items."""
+    if not _embedder:
+        logger.error("Cannot build FAISS index: Embedder not available.")
+        return None
+    
+    index = faiss.IndexFlatL2(_dimension)
+    if not items_list:
+        return index
+
+    logger.info(f"Building FAISS index for {len(items_list)} items...")
+    texts_to_embed = [text_extraction_fn(item) for item in items_list]
+    
+    try:
+        embeddings = _embedder.encode(texts_to_embed, convert_to_tensor=False, show_progress_bar=False)
+        embeddings_np = np.array(embeddings, dtype=np.float32)
+        if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(items_list):
+            index.add(embeddings_np)
+            logger.info(f"FAISS index built successfully with {index.ntotal} items.")
+        else:
+            logger.error(f"FAISS build failed: Embeddings shape error. Expected ({len(items_list)}, {_dimension}), Got {getattr(embeddings_np, 'shape', 'N/A')}")
+            return faiss.IndexFlatL2(_dimension) # Return empty index on failure
+    except Exception as e:
+        logger.error(f"Error building FAISS index: {e}", exc_info=True)
+        return faiss.IndexFlatL2(_dimension) # Return empty index on failure
+
+    return index
+
+# --- Initialization ---
 def initialize_memory_system():
     global _initialized, _embedder, _dimension, _faiss_memory_index, _memory_items_list, _faiss_rules_index, _rules_items_list
     
     with _init_lock:
         if _initialized:
-            logger.info("Memory system already initialized.")
             return
 
         logger.info(f"Initializing memory system with backend: {STORAGE_BACKEND}")
         init_start_time = time.time()
 
         if not SentenceTransformer or not faiss or not np:
-            logger.error("Core RAG libraries (SentenceTransformers, FAISS, NumPy) not available. Cannot initialize semantic memory.")
-            _initialized = False
+            logger.error("Core RAG libraries not available. Cannot initialize semantic memory.")
             return
         
         if not _embedder:
@@ -110,246 +137,185 @@ def initialize_memory_system():
                 logger.info("Loading SentenceTransformer model (all-MiniLM-L6-v2)...")
                 _embedder = SentenceTransformer('all-MiniLM-L6-v2', cache_folder="./sentence_transformer_cache")
                 _dimension = _embedder.get_sentence_embedding_dimension() or 384
-                logger.info(f"SentenceTransformer loaded. Dimension: {_dimension}")
             except Exception as e:
                 logger.critical(f"FATAL: Error loading SentenceTransformer: {e}", exc_info=True)
-                _initialized = False
                 return
 
-        if STORAGE_BACKEND == "SQLITE":
-            _init_sqlite_tables()
+        if STORAGE_BACKEND == "SQLITE": _init_sqlite_tables()
 
-        # Load Memories
-        logger.info("Loading memories from persistent storage...")
+        # Load Memories from persistent storage
         temp_memories_json = []
-        if STORAGE_BACKEND == "RAM":
-            pass
-        elif STORAGE_BACKEND == "SQLITE" and sqlite3:
-            try:
-                with _get_sqlite_connection() as conn:
-                    temp_memories_json = [row[0] for row in conn.execute("SELECT memory_json FROM memories ORDER BY created_at ASC")]
+        if STORAGE_BACKEND == "SQLITE":
+            try: temp_memories_json = [row[0] for row in _get_sqlite_connection().execute("SELECT memory_json FROM memories")]
             except Exception as e: logger.error(f"Error loading memories from SQLite: {e}")
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset and load_dataset:
+        elif STORAGE_BACKEND == "HF_DATASET":
             try:
-                logger.info(f"Attempting to load memories from HF Dataset: {HF_MEMORY_DATASET_REPO}")
+                logger.info(f"Loading memories from HF Dataset: {HF_MEMORY_DATASET_REPO}")
                 dataset = load_dataset(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
                 if "train" in dataset and "memory_json" in dataset["train"].column_names:
-                    num_rows = len(dataset["train"])
-                    logger.info(f"HF Dataset for memories found. 'train' split has {num_rows} rows.")
-                    temp_memories_json = [m_json for m_json in dataset["train"]["memory_json"] if isinstance(m_json, str) and m_json.strip()]
-                    logger.info(f"Extracted {len(temp_memories_json)} valid memory JSON strings from the dataset.")
-                else:
-                    logger.warning(f"HF Dataset {HF_MEMORY_DATASET_REPO} loaded, but 'train' split or 'memory_json' column is missing. Dataset structure: {dataset}")
-            except Exception as e: logger.error(f"Error loading memories from HF Dataset ({HF_MEMORY_DATASET_REPO}): {e}", exc_info=True)
+                    temp_memories_json = [m for m in dataset["train"]["memory_json"] if isinstance(m, str) and m.strip()]
+                    logger.info(f"Loaded {len(temp_memories_json)} valid memories from HF Dataset.")
+                else: logger.warning(f"HF Dataset {HF_MEMORY_DATASET_REPO} has no 'train' split or 'memory_json' column.")
+            except Exception as e: logger.error(f"Error loading memories from HF Dataset: {e}", exc_info=True)
         
         _memory_items_list = temp_memories_json
-        logger.info(f"Loaded {len(_memory_items_list)} memory items into cache from {STORAGE_BACKEND}.")
-
-        _faiss_memory_index = faiss.IndexFlatL2(_dimension)
-        if _memory_items_list:
-            logger.info(f"Building FAISS index for {len(_memory_items_list)} memories...")
-            texts_to_embed_mem = []
-            for mem_json_str in _memory_items_list:
-                try:
-                    mem_obj = json.loads(mem_json_str)
-                    text = f"User: {mem_obj.get('user_input','')}\nAI: {mem_obj.get('bot_response','')}\nTakeaway: {mem_obj.get('metrics',{}).get('takeaway','N/A')}"
-                    texts_to_embed_mem.append(text)
-                except json.JSONDecodeError: logger.warning(f"Skipping malformed memory JSON for FAISS indexing: {mem_json_str[:100]}")
-            
-            if texts_to_embed_mem:
-                try:
-                    embeddings = _embedder.encode(texts_to_embed_mem, convert_to_tensor=False, show_progress_bar=False)
-                    embeddings_np = np.array(embeddings, dtype=np.float32)
-                    if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(texts_to_embed_mem) and embeddings_np.shape[1] == _dimension:
-                        _faiss_memory_index.add(embeddings_np)
-                    else: logger.error(f"Memory embeddings shape error. Expected ({len(texts_to_embed_mem)}, {_dimension}), Got {embeddings_np.shape if hasattr(embeddings_np, 'shape') else 'N/A'}")
-                except Exception as e_faiss_mem: logger.error(f"Error building FAISS memory index: {e_faiss_mem}")
-        logger.info(f"FAISS memory index built. Total items: {_faiss_memory_index.ntotal if _faiss_memory_index else 'N/A'}")
-
-        # Load Rules
-        logger.info("Loading rules from persistent storage...")
+        
+        # Build Memory FAISS Index
+        _faiss_memory_index = _build_faiss_index(
+            _memory_items_list, 
+            lambda m: f"User: {json.loads(m).get('user_input', '')}\nAI: {json.loads(m).get('bot_response', '')}\nTakeaway: {json.loads(m).get('metrics', {}).get('takeaway', 'N/A')}"
+        )
+
+        # Load Rules from persistent storage
         temp_rules_text = []
-        if STORAGE_BACKEND == "RAM":
-            pass
-        elif STORAGE_BACKEND == "SQLITE" and sqlite3:
-            try:
-                with _get_sqlite_connection() as conn:
-                    temp_rules_text = [row[0] for row in conn.execute("SELECT rule_text FROM rules ORDER BY created_at ASC")]
+        if STORAGE_BACKEND == "SQLITE":
+            try: temp_rules_text = [row[0] for row in _get_sqlite_connection().execute("SELECT rule_text FROM rules")]
             except Exception as e: logger.error(f"Error loading rules from SQLite: {e}")
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset and load_dataset:
+        elif STORAGE_BACKEND == "HF_DATASET":
             try:
-                logger.info(f"Attempting to load rules from HF Dataset: {HF_RULES_DATASET_REPO}")
+                logger.info(f"Loading rules from HF Dataset: {HF_RULES_DATASET_REPO}")
                 dataset = load_dataset(HF_RULES_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
                 if "train" in dataset and "rule_text" in dataset["train"].column_names:
-                    num_rows = len(dataset["train"])
-                    logger.info(f"HF Dataset for rules found. 'train' split has {num_rows} rows.")
-                    temp_rules_text = [r_text for r_text in dataset["train"]["rule_text"] if isinstance(r_text, str) and r_text.strip()]
-                    logger.info(f"Extracted {len(temp_rules_text)} valid rule strings from the dataset.")
-                else:
-                    logger.warning(f"HF Dataset {HF_RULES_DATASET_REPO} for rules loaded, but 'train' split or 'rule_text' column is missing. Dataset structure: {dataset}")
-            except Exception as e: logger.error(f"Error loading rules from HF Dataset ({HF_RULES_DATASET_REPO}): {e}", exc_info=True)
-        
+                    temp_rules_text = [r for r in dataset["train"]["rule_text"] if isinstance(r, str) and r.strip()]
+                    logger.info(f"Loaded {len(temp_rules_text)} valid rules from HF Dataset.")
+                else: logger.warning(f"HF Dataset {HF_RULES_DATASET_REPO} has no 'train' split or 'rule_text' column.")
+            except Exception as e: logger.error(f"Error loading rules from HF Dataset: {e}", exc_info=True)
+
         _rules_items_list = sorted(list(set(temp_rules_text)))
-        logger.info(f"Loaded {len(_rules_items_list)} rule items into cache from {STORAGE_BACKEND}.")
-
-        _faiss_rules_index = faiss.IndexFlatL2(_dimension)
-        if _rules_items_list:
-            logger.info(f"Building FAISS index for {len(_rules_items_list)} rules...")
-            if _rules_items_list:
-                try:
-                    embeddings = _embedder.encode(_rules_items_list, convert_to_tensor=False, show_progress_bar=False)
-                    embeddings_np = np.array(embeddings, dtype=np.float32)
-                    if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(_rules_items_list) and embeddings_np.shape[1] == _dimension:
-                        _faiss_rules_index.add(embeddings_np)
-                    else: logger.error(f"Rule embeddings shape error. Expected ({len(_rules_items_list)}, {_dimension}), Got {embeddings_np.shape if hasattr(embeddings_np, 'shape') else 'N/A'}")
-                except Exception as e_faiss_rule: logger.error(f"Error building FAISS rule index: {e_faiss_rule}")
-        logger.info(f"FAISS rules index built. Total items: {_faiss_rules_index.ntotal if _faiss_rules_index else 'N/A'}")
+        
+        # Build Rules FAISS Index
+        _faiss_rules_index = _build_faiss_index(_rules_items_list, lambda r: r)
 
         _initialized = True
         logger.info(f"Memory system initialization complete in {time.time() - init_start_time:.2f}s")
 
-def _ensure_initialized(item_list, storage_type):
-    """Internal helper to check for initialization and trigger a reload if cache is empty on a persistent backend."""
-    global _initialized
-    if not _initialized or (storage_type != "RAM" and not item_list):
-        if not _initialized:
-             logger.warning("Memory system not initialized. Forcing initialization.")
-        else:
-             logger.warning(f"Persistent backend ({storage_type}) is configured, but cache is empty. Forcing re-initialization to reload data.")
-        
-        with _init_lock:
-             _initialized = False
-             initialize_memory_system()
+
+def _verify_and_rebuild_if_needed(index, items_list, text_extraction_fn):
+    """Self-healing function to ensure FAISS index is synced with the item list."""
+    if not index or index.ntotal != len(items_list):
+        logger.warning(
+            f"FAISS index mismatch detected (Index: {index.ntotal if index else 'None'}, List: {len(items_list)}). "
+            "Rebuilding index from in-memory cache."
+        )
+        return _build_faiss_index(items_list, text_extraction_fn)
+    return index
+
 
 # --- Memory Operations (Semantic) ---
 def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
     global _memory_items_list, _faiss_memory_index
-    _ensure_initialized(_memory_items_list, STORAGE_BACKEND)
+    if not _initialized: initialize_memory_system()
     if not _embedder or not _faiss_memory_index:
-        return False, "Memory system or embedder not initialized for adding memory."
+        return False, "Memory system not ready for adding entries."
+
     memory_obj = {"user_input": user_input, "metrics": metrics, "bot_response": bot_response, "timestamp": datetime.utcnow().isoformat()}
     memory_json_str = json.dumps(memory_obj)
     text_to_embed = f"User: {user_input}\nAI: {bot_response}\nTakeaway: {metrics.get('takeaway', 'N/A')}"
+    
     try:
         embedding = _embedder.encode([text_to_embed], convert_to_tensor=False)
-        embedding_np = np.array(embedding, dtype=np.float32).reshape(1, -1)
-        if embedding_np.shape != (1, _dimension):
-            logger.error(f"Memory embedding shape error: {embedding_np.shape}. Expected (1, {_dimension})")
-            return False, "Embedding shape error."
+        embedding_np = np.array(embedding, dtype=np.float32)
+        
         _faiss_memory_index.add(embedding_np)
         _memory_items_list.append(memory_json_str)
-        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+        
+        if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn:
-                conn.execute("INSERT INTO memories (memory_json) VALUES (?)", (memory_json_str,))
-                conn.commit()
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
-            logger.info(f"Pushing {len(_memory_items_list)} memories to HF Hub: {HF_MEMORY_DATASET_REPO}")
+                conn.execute("INSERT INTO memories (memory_json) VALUES (?)", (memory_json_str,)); conn.commit()
+        elif STORAGE_BACKEND == "HF_DATASET":
             Dataset.from_dict({"memory_json": list(_memory_items_list)}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
-        logger.info(f"Added memory. RAM: {len(_memory_items_list)}, FAISS: {_faiss_memory_index.ntotal}")
+        
+        logger.info(f"Added memory. Cache size: {len(_memory_items_list)}, FAISS size: {_faiss_memory_index.ntotal}")
         return True, "Memory added successfully."
     except Exception as e:
         logger.error(f"Error adding memory entry: {e}", exc_info=True)
         return False, f"Error adding memory: {e}"
 
 def retrieve_memories_semantic(query: str, k: int = 3) -> list[dict]:
-    _ensure_initialized(_memory_items_list, STORAGE_BACKEND)
-    if not _embedder or not _faiss_memory_index or _faiss_memory_index.ntotal == 0:
-        logger.warning("Cannot retrieve memories: Embedder/FAISS index not ready or empty after initialization attempt.")
+    global _faiss_memory_index
+    if not _initialized: initialize_memory_system()
+    
+    # Self-healing: Verify index is synced with cache, rebuild if not.
+    _faiss_memory_index = _verify_and_rebuild_if_needed(
+        _faiss_memory_index, _memory_items_list,
+        lambda m: f"User: {json.loads(m).get('user_input', '')}\nAI: {json.loads(m).get('bot_response', '')}\nTakeaway: {json.loads(m).get('metrics', {}).get('takeaway', 'N/A')}"
+    )
+
+    if not _faiss_memory_index or _faiss_memory_index.ntotal == 0:
+        logger.debug("Cannot retrieve memories: index is empty.")
         return []
+
     try:
         query_embedding = _embedder.encode([query], convert_to_tensor=False)
-        query_embedding_np = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
-        if query_embedding_np.shape[1] != _dimension:
-            logger.error(f"Query embedding dimension mismatch. Expected {_dimension}, got {query_embedding_np.shape[1]}")
-            return []
+        query_embedding_np = np.array(query_embedding, dtype=np.float32)
         distances, indices = _faiss_memory_index.search(query_embedding_np, min(k, _faiss_memory_index.ntotal))
-        results = []
-        for i in indices[0]:
-            if 0 <= i < len(_memory_items_list):
-                try: results.append(json.loads(_memory_items_list[i]))
-                except json.JSONDecodeError: logger.warning(f"Could not parse memory JSON from list at index {i}")
-            else: logger.warning(f"FAISS index {i} out of bounds for memory_items_list (len: {len(_memory_items_list)})")
-        logger.debug(f"Retrieved {len(results)} memories semantically for query: '{query[:50]}...'")
+        
+        results = [json.loads(_memory_items_list[i]) for i in indices[0] if 0 <= i < len(_memory_items_list)]
+        logger.info(f"Retrieved {len(results)} memories for query: '{query[:50]}...'")
         return results
     except Exception as e:
         logger.error(f"Error retrieving memories semantically: {e}", exc_info=True)
         return []
 
+
 # --- Rule (Insight) Operations (Semantic) ---
 def add_rule_entry(rule_text: str) -> tuple[bool, str]:
     global _rules_items_list, _faiss_rules_index
-    _ensure_initialized(_rules_items_list, STORAGE_BACKEND)
-    if not _embedder or not _faiss_rules_index: return False, "Rule system or embedder not initialized."
+    if not _initialized: initialize_memory_system()
+
     rule_text = rule_text.strip()
-    if not rule_text: return False, "Rule text cannot be empty."
-    if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\](.*)", rule_text, re.I|re.DOTALL):
+    if not rule_text or "duplicate" == rule_text or rule_text in _rules_items_list:
+        return False, "duplicate or invalid"
+    if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", rule_text, re.I):
         return False, "Invalid rule format."
-    if rule_text in _rules_items_list: return False, "duplicate"
+
     try:
         embedding = _embedder.encode([rule_text], convert_to_tensor=False)
-        embedding_np = np.array(embedding, dtype=np.float32).reshape(1, -1)
-        if embedding_np.shape != (1, _dimension): return False, "Rule embedding shape error."
+        embedding_np = np.array(embedding, dtype=np.float32)
         _faiss_rules_index.add(embedding_np)
         _rules_items_list.append(rule_text)
         _rules_items_list.sort()
-        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+
+        if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn:
-                conn.execute("INSERT OR IGNORE INTO rules (rule_text) VALUES (?)", (rule_text,))
-                conn.commit()
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
-            logger.info(f"Pushing {len(_rules_items_list)} rules to HF Hub: {HF_RULES_DATASET_REPO}")
+                conn.execute("INSERT OR IGNORE INTO rules (rule_text) VALUES (?)", (rule_text,)); conn.commit()
+        elif STORAGE_BACKEND == "HF_DATASET":
             Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
-        logger.info(f"Added rule. RAM: {len(_rules_items_list)}, FAISS: {_faiss_rules_index.ntotal}")
         return True, "Rule added successfully."
     except Exception as e:
         logger.error(f"Error adding rule entry: {e}", exc_info=True)
         return False, f"Error adding rule: {e}"
 
 def retrieve_rules_semantic(query: str, k: int = 5) -> list[str]:
-    _ensure_initialized(_rules_items_list, STORAGE_BACKEND)
-    if not _embedder or not _faiss_rules_index or _faiss_rules_index.ntotal == 0:
-        logger.warning("Cannot retrieve rules: Embedder/FAISS index not ready or empty after initialization attempt.")
-        return []
+    global _faiss_rules_index
+    if not _initialized: initialize_memory_system()
+
+    _faiss_rules_index = _verify_and_rebuild_if_needed(_faiss_rules_index, _rules_items_list, lambda r: r)
+
+    if not _faiss_rules_index or _faiss_rules_index.ntotal == 0: return []
     try:
         query_embedding = _embedder.encode([query], convert_to_tensor=False)
-        query_embedding_np = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
-        if query_embedding_np.shape[1] != _dimension: return []
+        query_embedding_np = np.array(query_embedding, dtype=np.float32)
         distances, indices = _faiss_rules_index.search(query_embedding_np, min(k, _faiss_rules_index.ntotal))
-        results = [_rules_items_list[i] for i in indices[0] if 0 <= i < len(_rules_items_list)]
-        logger.debug(f"Retrieved {len(results)} rules semantically for query: '{query[:50]}...'")
-        return results
+        return [_rules_items_list[i] for i in indices[0] if 0 <= i < len(_rules_items_list)]
     except Exception as e:
         logger.error(f"Error retrieving rules semantically: {e}", exc_info=True)
         return []
 
 def remove_rule_entry(rule_text_to_delete: str) -> bool:
     global _rules_items_list, _faiss_rules_index
-    _ensure_initialized(_rules_items_list, STORAGE_BACKEND)
-    if not _embedder or not _faiss_rules_index: return False
+    if not _initialized: initialize_memory_system()
     rule_text_to_delete = rule_text_to_delete.strip()
     if rule_text_to_delete not in _rules_items_list: return False
     try:
         _rules_items_list.remove(rule_text_to_delete)
-        _rules_items_list.sort()
-        new_faiss_rules_index = faiss.IndexFlatL2(_dimension)
-        if _rules_items_list:
-            embeddings = _embedder.encode(_rules_items_list, convert_to_tensor=False)
-            embeddings_np = np.array(embeddings, dtype=np.float32)
-            if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(_rules_items_list) and embeddings_np.shape[1] == _dimension:
-                 new_faiss_rules_index.add(embeddings_np)
-            else:
-                logger.error("Error rebuilding FAISS for rules after removal: Embedding shape error. State might be inconsistent.")
-                _rules_items_list.append(rule_text_to_delete)
-                _rules_items_list.sort()
-                return False
-        _faiss_rules_index = new_faiss_rules_index
-        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+        _faiss_rules_index = _build_faiss_index(_rules_items_list, lambda r: r)
+        
+        if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn:
-                conn.execute("DELETE FROM rules WHERE rule_text = ?", (rule_text_to_delete,))
-                conn.commit()
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
+                conn.execute("DELETE FROM rules WHERE rule_text = ?", (rule_text_to_delete,)); conn.commit()
+        elif STORAGE_BACKEND == "HF_DATASET":
             Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
-        logger.info(f"Removed rule. RAM: {len(_rules_items_list)}, FAISS: {_faiss_rules_index.ntotal}")
         return True
     except Exception as e:
         logger.error(f"Error removing rule entry: {e}", exc_info=True)
@@ -357,50 +323,42 @@ def remove_rule_entry(rule_text_to_delete: str) -> bool:
 
 # --- Utility functions to get all data (for UI display, etc.) ---
 def get_all_rules_cached() -> list[str]:
-    _ensure_initialized(_rules_items_list, STORAGE_BACKEND)
+    if not _initialized: initialize_memory_system()
     return list(_rules_items_list)
 
 def get_all_memories_cached() -> list[dict]:
-    _ensure_initialized(_memory_items_list, STORAGE_BACKEND)
-    mem_dicts = []
-    for mem_json_str in _memory_items_list:
-        try: mem_dicts.append(json.loads(mem_json_str))
-        except: pass
-    return mem_dicts
+    if not _initialized: initialize_memory_system()
+    return [json.loads(m) for m in _memory_items_list if m]
 
 def clear_all_memory_data_backend() -> bool:
     global _memory_items_list, _faiss_memory_index
-    _ensure_initialized(_memory_items_list, STORAGE_BACKEND)
-    success = True
+    if not _initialized: initialize_memory_system()
+    _memory_items_list = []
+    if _faiss_memory_index: _faiss_memory_index.reset()
     try:
-        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+        if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn: conn.execute("DELETE FROM memories"); conn.commit()
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
+        elif STORAGE_BACKEND == "HF_DATASET":
             Dataset.from_dict({"memory_json": []}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
-        _memory_items_list = []
-        if _faiss_memory_index: _faiss_memory_index.reset()
-        logger.info("All memories cleared from backend and in-memory stores.")
+        logger.info("All memories cleared.")
+        return True
     except Exception as e:
-        logger.error(f"Error clearing all memory data: {e}")
-        success = False
-    return success
+        logger.error(f"Error clearing all memory data: {e}"); return False
 
 def clear_all_rules_data_backend() -> bool:
     global _rules_items_list, _faiss_rules_index
-    _ensure_initialized(_rules_items_list, STORAGE_BACKEND)
-    success = True
+    if not _initialized: initialize_memory_system()
+    _rules_items_list = []
+    if _faiss_rules_index: _faiss_rules_index.reset()
     try:
-        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+        if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn: conn.execute("DELETE FROM rules"); conn.commit()
-        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
+        elif STORAGE_BACKEND == "HF_DATASET":
             Dataset.from_dict({"rule_text": []}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
-        _rules_items_list = []
-        if _faiss_rules_index: _faiss_rules_index.reset()
-        logger.info("All rules cleared from backend and in-memory stores.")
+        logger.info("All rules cleared.")
+        return True
     except Exception as e:
-        logger.error(f"Error clearing all rules data: {e}")
-        success = False
-    return success
+        logger.error(f"Error clearing all rules data: {e}"); return False
 
 FAISS_MEMORY_PATH = os.path.join(os.getenv("FAISS_STORAGE_PATH", "app_data/faiss_indices"), "memory_index.faiss")
 FAISS_RULES_PATH = os.path.join(os.getenv("FAISS_STORAGE_PATH", "app_data/faiss_indices"), "rules_index.faiss")
@@ -410,32 +368,14 @@ def save_faiss_indices_to_disk():
     faiss_dir = os.path.dirname(FAISS_MEMORY_PATH)
     if not os.path.exists(faiss_dir): os.makedirs(faiss_dir, exist_ok=True)
     if _faiss_memory_index and _faiss_memory_index.ntotal > 0:
-        try:
-            faiss.write_index(_faiss_memory_index, FAISS_MEMORY_PATH)
-            logger.info(f"Memory FAISS index saved to disk ({_faiss_memory_index.ntotal} items).")
-        except Exception as e: logger.error(f"Error saving memory FAISS index: {e}")
+        faiss.write_index(_faiss_memory_index, FAISS_MEMORY_PATH)
     if _faiss_rules_index and _faiss_rules_index.ntotal > 0:
-        try:
-            faiss.write_index(_faiss_rules_index, FAISS_RULES_PATH)
-            logger.info(f"Rules FAISS index saved to disk ({_faiss_rules_index.ntotal} items).")
-        except Exception as e: logger.error(f"Error saving rules FAISS index: {e}")
+        faiss.write_index(_faiss_rules_index, FAISS_RULES_PATH)
 
 def load_faiss_indices_from_disk():
     global _faiss_memory_index, _faiss_rules_index
     if not _initialized or not faiss: return
-    if os.path.exists(FAISS_MEMORY_PATH) and _faiss_memory_index:
-        try:
-            logger.info(f"Loading memory FAISS index from {FAISS_MEMORY_PATH}...")
-            _faiss_memory_index = faiss.read_index(FAISS_MEMORY_PATH)
-            logger.info(f"Memory FAISS index loaded ({_faiss_memory_index.ntotal} items).")
-            if _faiss_memory_index.ntotal != len(_memory_items_list) and len(_memory_items_list) > 0:
-                 logger.warning(f"Memory FAISS index count ({_faiss_memory_index.ntotal}) differs from loaded texts ({len(_memory_items_list)}). Consider rebuilding FAISS.")
-        except Exception as e: logger.error(f"Error loading memory FAISS index: {e}. Will use fresh index.")
-    if os.path.exists(FAISS_RULES_PATH) and _faiss_rules_index:
-        try:
-            logger.info(f"Loading rules FAISS index from {FAISS_RULES_PATH}...")
-            _faiss_rules_index = faiss.read_index(FAISS_RULES_PATH)
-            logger.info(f"Rules FAISS index loaded ({_faiss_rules_index.ntotal} items).")
-            if _faiss_rules_index.ntotal != len(_rules_items_list) and len(_rules_items_list) > 0:
-                 logger.warning(f"Rules FAISS index count ({_faiss_rules_index.ntotal}) differs from loaded texts ({len(_rules_items_list)}). Consider rebuilding FAISS.")
-        except Exception as e: logger.error(f"Error loading rules FAISS index: {e}. Will use fresh index.")
+    if os.path.exists(FAISS_MEMORY_PATH):
+        _faiss_memory_index = faiss.read_index(FAISS_MEMORY_PATH)
+    if os.path.exists(FAISS_RULES_PATH):
+        _faiss_rules_index = faiss.read_index(FAISS_RULES_PATH)