Update memory_logic.py

memory_logic.py

@@ -1,4 +1,3 @@
-# memory_logic.py
 import os
 import json
 import time
@@ -7,7 +6,6 @@ import logging
 import re
 import threading
 
-# Conditionally import heavy dependencies
 try:
     from sentence_transformers import SentenceTransformer
     import faiss
@@ -28,21 +26,16 @@ except ImportError:
     load_dataset, Dataset = None, None
     logging.warning("datasets library not installed. Hugging Face Dataset backend will be unavailable.")
 
-
 logger = logging.getLogger(__name__)
-# Suppress verbose logs from dependencies
 for lib_name in ["sentence_transformers", "faiss", "datasets", "huggingface_hub"]:
     if logging.getLogger(lib_name): logging.getLogger(lib_name).setLevel(logging.WARNING)
 
-
-# --- Configuration (Read directly from environment variables) ---
-STORAGE_BACKEND = os.getenv("STORAGE_BACKEND", "HF_DATASET").upper() #HF_DATASET, RAM, SQLITE
+STORAGE_BACKEND = os.getenv("STORAGE_BACKEND", "HF_DATASET").upper()
 SQLITE_DB_PATH = os.getenv("SQLITE_DB_PATH", "app_data/ai_memory.db")
 HF_TOKEN = os.getenv("HF_TOKEN")
 HF_MEMORY_DATASET_REPO = os.getenv("HF_MEMORY_DATASET_REPO", "broadfield-dev/ai-brain")
 HF_RULES_DATASET_REPO = os.getenv("HF_RULES_DATASET_REPO", "broadfield-dev/ai-rules")
 
-# --- Globals for RAG within this module ---
 _embedder = None
 _dimension = 384
 _faiss_memory_index = None
@@ -74,10 +67,10 @@ def _init_sqlite_tables():
 def _build_faiss_index(items_list, text_extraction_fn):
     if not _embedder:
         logger.error("Cannot build FAISS index: Embedder not available.")
-        return None
+        return None, []
     
     index = faiss.IndexFlatL2(_dimension)
-    if not items_list: return index
+    if not items_list: return index, []
 
     texts_to_embed, valid_items = [], []
     for item in items_list:
@@ -89,7 +82,7 @@ def _build_faiss_index(items_list, text_extraction_fn):
 
     if not texts_to_embed:
         logger.warning("No valid items to embed for FAISS index after filtering.")
-        return index
+        return index, []
 
     try:
         embeddings = _embedder.encode(texts_to_embed, convert_to_tensor=False, show_progress_bar=False)
@@ -97,7 +90,6 @@ def _build_faiss_index(items_list, text_extraction_fn):
         if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(texts_to_embed):
             index.add(embeddings_np)
             logger.info(f"FAISS index built with {index.ntotal} / {len(items_list)} items.")
-            # Important: The original items_list is returned so we can update the global list to match the index
             return index, valid_items
         else:
             logger.error(f"FAISS build failed: Embeddings shape error.")
@@ -126,11 +118,10 @@ def initialize_memory_system():
                 _dimension = _embedder.get_sentence_embedding_dimension() or 384
             except Exception as e:
                 logger.critical(f"FATAL: Could not load SentenceTransformer model. Semantic search disabled. Error: {e}", exc_info=True)
-                return # Stop initialization if embedder fails
+                return
 
         if STORAGE_BACKEND == "SQLITE": _init_sqlite_tables()
 
-        # Load raw data
         raw_mems = []
         if STORAGE_BACKEND == "SQLITE":
             try: raw_mems = [row[0] for row in _get_sqlite_connection().execute("SELECT memory_json FROM memories")]
@@ -142,13 +133,11 @@ def initialize_memory_system():
                     raw_mems = [m for m in dataset["train"]["memory_json"] if isinstance(m, str) and m.strip()]
             except Exception as e: logger.error(f"Error loading memories from HF Dataset: {e}", exc_info=True)
         
-        # Build Memory Index and get validated list
         mem_index, valid_mems = _build_faiss_index(raw_mems, lambda m: f"User: {json.loads(m).get('user_input', '')}\nAI: {json.loads(m).get('bot_response', '')}")
         _faiss_memory_index = mem_index
-        _memory_items_list = valid_mems # Use the validated list that matches the index
+        _memory_items_list = valid_mems
         logger.info(f"Loaded and indexed {len(_memory_items_list)} memories.")
 
-        # Load Rules
         raw_rules = []
         if STORAGE_BACKEND == "SQLITE":
             try: raw_rules = [row[0] for row in _get_sqlite_connection().execute("SELECT rule_text FROM rules")]
@@ -160,13 +149,11 @@ 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)
 
-        # Build Rules Index and get validated list
         rule_index, valid_rules = _build_faiss_index(sorted(list(set(raw_rules))), lambda r: r)
         _faiss_rules_index = rule_index
-        _rules_items_list = valid_rules # Use the validated list that matches the index
+        _rules_items_list = valid_rules
         logger.info(f"Loaded and indexed {len(_rules_items_list)} rules.")
         
-        # Only mark as initialized if the core components are ready
         if _embedder and _faiss_memory_index is not None and _faiss_rules_index is not None:
             _initialized = True
             logger.info(f"Memory system initialization complete in {time.time() - init_start_time:.2f}s")
@@ -174,18 +161,19 @@ def initialize_memory_system():
             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)
-        # This function is now stateful and modifies the global list to match the new index
-        if isinstance(items_list, list) and isinstance(valid_items, list):
-             # A bit of a hack to update the global list from here
-             items_list.clear()
-             items_list.extend(valid_items)
+        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()
     if not _embedder or _faiss_memory_index is None: return False, "Memory system not ready."
 
@@ -225,6 +213,7 @@ def retrieve_memories_semantic(query: str, k: int = 3) -> list[dict]:
         return []
 
 def add_rule_entry(rule_text: str) -> tuple[bool, str]:
+    global _rules_items_list, _faiss_rules_index
     if not _initialized: initialize_memory_system()
     if not _embedder or _faiss_rules_index is None: return False, "Rule system not ready."
 
@@ -264,11 +253,11 @@ def retrieve_rules_semantic(query: str, k: int = 5) -> list[str]:
         return []
 
 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:
-        # Rebuild index and list without the deleted rule
         new_list = [r for r in _rules_items_list if r != rule_text_to_delete]
         _faiss_rules_index, _rules_items_list = _build_faiss_index(new_list, lambda r: r)
         
@@ -292,7 +281,7 @@ def get_all_memories_cached() -> list[dict]:
         try:
             valid_mems.append(json.loads(m_str))
         except json.JSONDecodeError:
-            continue # Skip corrupted data for UI display
+            continue
     return valid_mems
 
 def clear_all_memory_data_backend() -> bool:
@@ -324,8 +313,6 @@ def clear_all_rules_data_backend() -> bool:
         logger.error(f"Error clearing rules data: {e}"); return False
 
 def save_faiss_indices_to_disk():
-    # This function is primarily for the RAM backend, which is not the focus here.
-    # It's kept for compatibility.
     if not _initialized or not faiss: return
     faiss_dir = "app_data/faiss_indices"
     os.makedirs(faiss_dir, exist_ok=True)