Update memory_logic.py

memory_logic.py

@@ -149,7 +149,7 @@ def initialize_memory_system():
                     raw_rules = [r for r in dataset["train"]["rule_text"] if isinstance(r, str) and r.strip()]
             except Exception as e: logger.error(f"Error loading rules from HF Dataset: {e}", exc_info=True)
 
-        rule_index, valid_rules = _build_faiss_index(sorted(list(set(raw_rules))), lambda r: r)
+        rule_index, valid_rules = _build_faiss_index(raw_rules, lambda r: r)
         _faiss_rules_index = rule_index
         _rules_items_list = valid_rules
         logger.info(f"Loaded and indexed {len(_rules_items_list)} rules.")
@@ -160,18 +160,6 @@ def initialize_memory_system():
         else:
             logger.error("Memory system initialization failed. Core components are not ready.")
 
-def _verify_and_rebuild_if_needed(index, items_list, text_extraction_fn):
-    global _memory_items_list, _rules_items_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...")
-        new_index, valid_items = _build_faiss_index(items_list, text_extraction_fn)
-        if items_list is _memory_items_list:
-            _memory_items_list = valid_items
-        elif items_list is _rules_items_list:
-            _rules_items_list = valid_items
-        return new_index
-    return index
-
 def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
     global _memory_items_list, _faiss_memory_index
     if not _initialized: initialize_memory_system()
@@ -197,12 +185,8 @@ def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple
         return False, f"Error: {e}"
 
 def retrieve_memories_semantic(query: str, k: int = 3) -> list[dict]:
-    global _faiss_memory_index
     if not _initialized: initialize_memory_system()
     if not _faiss_memory_index or _faiss_memory_index.ntotal == 0: return []
-
-    _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', '')}")
-    if not _faiss_memory_index or _faiss_memory_index.ntotal == 0: return []
     
     try:
         query_embedding = _embedder.encode([query], convert_to_tensor=False)
@@ -225,7 +209,6 @@ def add_rule_entry(rule_text: str) -> tuple[bool, str]:
         embedding = _embedder.encode([rule_text], convert_to_tensor=False)
         _faiss_rules_index.add(np.array(embedding, dtype=np.float32))
         _rules_items_list.append(rule_text)
-        _rules_items_list.sort()
 
         if STORAGE_BACKEND == "SQLITE":
             with _get_sqlite_connection() as conn: conn.execute("INSERT OR IGNORE INTO rules (rule_text) VALUES (?)", (rule_text,)); conn.commit()
@@ -237,13 +220,9 @@ def add_rule_entry(rule_text: str) -> tuple[bool, str]:
         return False, f"Error: {e}"
 
 def retrieve_rules_semantic(query: str, k: int = 5) -> list[str]:
-    global _faiss_rules_index
     if not _initialized: initialize_memory_system()
     if not _faiss_rules_index or _faiss_rules_index.ntotal == 0: return []
     
-    _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)
         distances, indices = _faiss_rules_index.search(np.array(query_embedding, dtype=np.float32), min(k, _faiss_rules_index.ntotal))
@@ -256,17 +235,23 @@ def remove_rule_entry(rule_text_to_delete: str) -> bool:
     global _rules_items_list, _faiss_rules_index
     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:
-        new_list = [r for r in _rules_items_list if r != rule_text_to_delete]
-        new_index, valid_items = _build_faiss_index(new_list, lambda r: r)
-        _faiss_rules_index = new_index
-        _rules_items_list = valid_items
+        idx_to_remove = _rules_items_list.index(rule_text_to_delete)
+    except ValueError:
+        return False
+
+    try:
+        _faiss_rules_index.remove_ids(np.array([idx_to_remove], dtype='int64'))
+        del _rules_items_list[idx_to_remove]
         
         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 _rules_items_list:
+        elif STORAGE_BACKEND == "HF_DATASET":
+            # After removing, we need to push the new state of the list.
+            # Important: This can be slow if the dataset is large.
             Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
+        
         return True
     except Exception as e:
         logger.error(f"Error removing rule: {e}", exc_info=True)
@@ -274,7 +259,7 @@ def remove_rule_entry(rule_text_to_delete: str) -> bool:
 
 def get_all_rules_cached() -> list[str]:
     if not _initialized: initialize_memory_system()
-    return list(_rules_items_list)
+    return sorted(list(_rules_items_list))
 
 def get_all_memories_cached() -> list[dict]:
     if not _initialized: initialize_memory_system()