Update memory_logic.py

memory_logic.py

@@ -1,4 +1,3 @@
-# memory_logic.py
 import os
 import json
 import time
@@ -60,18 +59,40 @@ def _get_sqlite_connection():
         os.makedirs(db_dir, exist_ok=True)
     return sqlite3.connect(SQLITE_DB_PATH, timeout=10)
 
+def _init_sqlite_tables():
+    if STORAGE_BACKEND != "SQLITE" or not sqlite3:
+        return
+    try:
+        with _get_sqlite_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS memories (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    memory_json TEXT NOT NULL,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS rules (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    rule_text TEXT NOT NULL UNIQUE,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            conn.commit()
+    except Exception as e:
+        logger.error(f"SQLite table initialization error: {e}", exc_info=True)
+
 def _build_faiss_index_from_json_strings(memory_items: list[str]) -> faiss.Index | None:
     if not memory_items or not _embedder:
         return faiss.IndexFlatL2(_dimension)
 
     texts_to_embed = []
-    valid_indices = []
-    for i, mem_json_str in enumerate(memory_items):
+    for mem_json_str in memory_items:
         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.append(text)
-            valid_indices.append(i)
         except json.JSONDecodeError:
             continue
     
@@ -103,8 +124,7 @@ def initialize_memory_system():
             return
 
         logger.info(f"Initializing memory system with backend: {STORAGE_BACKEND}")
-        init_start_time = time.time()
-
+        
         if not all([SentenceTransformer, faiss, np]):
             logger.error("Core RAG libraries not available. Cannot initialize semantic memory.")
             return
@@ -117,6 +137,9 @@ def initialize_memory_system():
                 logger.critical(f"FATAL: Error loading SentenceTransformer: {e}", exc_info=True)
                 return
 
+        if STORAGE_BACKEND == "SQLITE":
+            _init_sqlite_tables()
+            
         long_term_mems = []
         if STORAGE_BACKEND == "SQLITE" and sqlite3:
             try:
@@ -131,13 +154,10 @@ def initialize_memory_system():
             except Exception as e: logger.error(f"Error loading long-term memories from HF Dataset: {e}")
 
         _long_term_memory_items_list = long_term_mems
-        logger.info(f"Loaded {len(_long_term_memory_items_list)} long-term memory items.")
         _faiss_long_term_memory_index = _build_faiss_index_from_json_strings(_long_term_memory_items_list)
-        logger.info(f"Long-term memory FAISS index built. Total items: {_faiss_long_term_memory_index.ntotal if _faiss_long_term_memory_index else 'N/A'}")
-
+        
         _short_term_memory_items_list = []
         _faiss_short_term_memory_index = faiss.IndexFlatL2(_dimension)
-        logger.info("Short-term memory initialized (empty).")
         
         temp_rules_text = []
         if STORAGE_BACKEND == "SQLITE" and sqlite3:
@@ -156,10 +176,8 @@ def initialize_memory_system():
         if _rules_items_list:
             rule_embeddings = _embedder.encode(_rules_items_list, convert_to_tensor=False)
             _faiss_rules_index.add(np.array(rule_embeddings, dtype=np.float32))
-        logger.info(f"Rules FAISS index built. Total items: {_faiss_rules_index.ntotal if _faiss_rules_index else 'N/A'}")
 
         _initialized = True
-        logger.info(f"Memory system initialization complete in {time.time() - init_start_time:.2f}s")
 
 def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
     if not _initialized: initialize_memory_system()
@@ -184,7 +202,6 @@ def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple
             all_mems_for_push = _long_term_memory_items_list + _short_term_memory_items_list
             Dataset.from_dict({"memory_json": list(set(all_mems_for_push))}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
         
-        logger.info(f"Added memory. Short-term count: {_faiss_short_term_memory_index.ntotal}")
         return True, "Memory added successfully."
     except Exception as e:
         logger.error(f"Error adding memory entry: {e}", exc_info=True)
@@ -203,13 +220,11 @@ def search_memories(query: str, k: int = 3, threshold: float = 1.0) -> tuple[lis
         best_dist = distances[0][0] if len(distances[0]) > 0 else float('inf')
         
         if best_dist < threshold:
-            logger.info(f"Found relevant short-term memories (best distance: {best_dist:.4f}).")
             for i in indices[0]:
                 res = json.loads(_short_term_memory_items_list[i])
                 final_results[res['timestamp']] = res
             return list(final_results.values()), search_path
 
-    logger.info("No relevant short-term memories found. Escalating to deep search on long-term memory.")
     search_path = "deep"
     
     if _faiss_long_term_memory_index and _faiss_long_term_memory_index.ntotal > 0:
@@ -245,7 +260,6 @@ def get_all_memories_cached() -> list[dict]:
         except: continue
     return unique_mem_dicts
 
-# --- The rest of the utility functions (add_rule, get_rules, clear functions) remain the same ---
 def add_rule_entry(rule_text: str):
     global _rules_items_list, _faiss_rules_index
     if not _initialized: initialize_memory_system()
@@ -270,6 +284,74 @@ def add_rule_entry(rule_text: str):
         logger.error(f"Error adding rule: {e}", exc_info=True)
         return False, str(e)
 
+def remove_rule_entry(rule_text_to_delete: str) -> bool:
+    global _rules_items_list, _faiss_rules_index
+    if not _initialized: initialize_memory_system()
+    if not _embedder or not _faiss_rules_index: return False
+    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[1] == _dimension:
+                 new_faiss_rules_index.add(embeddings_np)
+            else:
+                _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:
+            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:
+            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 entry: {e}", exc_info=True)
+        return False
+
 def get_all_rules_cached() -> list[str]:
     if not _initialized: initialize_memory_system()
-    return list(_rules_items_list)
+    return list(_rules_items_list)
+
+def clear_all_memory_data_backend() -> bool:
+    global _long_term_memory_items_list, _short_term_memory_items_list, _faiss_long_term_memory_index, _faiss_short_term_memory_index
+    if not _initialized: initialize_memory_system()
+    success = True
+    try:
+        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+            with _get_sqlite_connection() as conn: conn.execute("DELETE FROM memories"); conn.commit()
+        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
+            Dataset.from_dict({"memory_json": []}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
+        _long_term_memory_items_list = []
+        _short_term_memory_items_list = []
+        if _faiss_long_term_memory_index: _faiss_long_term_memory_index.reset()
+        if _faiss_short_term_memory_index: _faiss_short_term_memory_index.reset()
+    except Exception as e:
+        logger.error(f"Error clearing all memory data: {e}")
+        success = False
+    return success
+
+def clear_all_rules_data_backend() -> bool:
+    global _rules_items_list, _faiss_rules_index
+    if not _initialized: initialize_memory_system()
+    success = True
+    try:
+        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+            with _get_sqlite_connection() as conn: conn.execute("DELETE FROM rules"); conn.commit()
+        elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and 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()
+    except Exception as e:
+        logger.error(f"Error clearing all rules data: {e}")
+        success = False
+    return success
+
+def save_faiss_indices_to_disk(): pass