memory_logic.py

# memory_logic.py
import os
import json
import time
from datetime import datetime
import logging
import re
import threading

# Conditionally import heavy dependencies
try:
    from sentence_transformers import SentenceTransformer
    import faiss
    import numpy as np
except ImportError:
    SentenceTransformer, faiss, np = None, None, None
    logging.warning("SentenceTransformers, FAISS, or NumPy not installed. Semantic search will be unavailable.")

try:
    import sqlite3
except ImportError:
    sqlite3 = None
    logging.warning("sqlite3 module not available. SQLite backend will be unavailable.")

try:
    from datasets import load_dataset, Dataset
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
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
_memory_items_list = []
_faiss_rules_index = None
_rules_items_list = []

_initialized = False
_init_lock = threading.Lock()

def _get_sqlite_connection():
    if not sqlite3: raise ImportError("sqlite3 module is required for SQLite backend.")
    db_dir = os.path.dirname(SQLITE_DB_PATH)
    if db_dir: 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, memory_json TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)")
            cursor.execute("CREATE TABLE IF NOT EXISTS rules (id INTEGER PRIMARY KEY, rule_text TEXT NOT NULL UNIQUE, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)")
            conn.commit()
        logger.info("SQLite tables checked/created.")
    except Exception as e:
        logger.error(f"SQLite table initialization error: {e}", exc_info=True)

def _build_faiss_index(items_list, text_extraction_fn):
    if not _embedder:
        logger.error("Cannot build FAISS index: Embedder not available.")
        return None
    
    index = faiss.IndexFlatL2(_dimension)
    if not items_list: return index

    texts_to_embed, valid_items = [], []
    for item in items_list:
        try:
            texts_to_embed.append(text_extraction_fn(item))
            valid_items.append(item)
        except (json.JSONDecodeError, TypeError) as e:
            logger.warning(f"Skipping item during FAISS indexing due to processing error: {e}. Item: '{str(item)[:100]}...'")

    if not texts_to_embed:
        logger.warning("No valid items to embed for FAISS index after filtering.")
        return index

    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(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.")
            return faiss.IndexFlatL2(_dimension), []
    except Exception as e:
        logger.error(f"Error building FAISS index: {e}", exc_info=True)
        return faiss.IndexFlatL2(_dimension), []

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: 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
        
        if not _embedder:
            try:
                logger.info("Loading SentenceTransformer model...")
                _embedder = SentenceTransformer('all-MiniLM-L6-v2', cache_folder="./sentence_transformer_cache")
                _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

        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")]
            except Exception as e: logger.error(f"Error loading memories from SQLite: {e}")
        elif STORAGE_BACKEND == "HF_DATASET":
            try:
                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:
                    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
        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")]
            except Exception as e: logger.error(f"Error loading rules from SQLite: {e}")
        elif STORAGE_BACKEND == "HF_DATASET":
            try:
                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:
                    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
        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")
        else:
            logger.error("Memory system initialization failed. Core components are not ready.")

def _verify_and_rebuild_if_needed(index, items_list, text_extraction_fn):
    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)
        return new_index
    return index

def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
    if not _initialized: initialize_memory_system()
    if not _embedder or _faiss_memory_index is None: return False, "Memory system not ready."

    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)
        _faiss_memory_index.add(np.array(embedding, dtype=np.float32))
        _memory_items_list.append(memory_json_str)
        
        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":
            Dataset.from_dict({"memory_json": list(_memory_items_list)}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
        
        return True, "Memory added."
    except Exception as e:
        logger.error(f"Error adding memory entry: {e}", exc_info=True)
        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)
        distances, indices = _faiss_memory_index.search(np.array(query_embedding, dtype=np.float32), min(k, _faiss_memory_index.ntotal))
        return [json.loads(_memory_items_list[i]) for i in indices[0] if 0 <= i < len(_memory_items_list)]
    except Exception as e:
        logger.error(f"Error retrieving memories: {e}", exc_info=True)
        return []

def add_rule_entry(rule_text: str) -> tuple[bool, str]:
    if not _initialized: initialize_memory_system()
    if not _embedder or _faiss_rules_index is None: return False, "Rule system not ready."

    rule_text = rule_text.strip()
    if not 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 format."

    try:
        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()
        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)
        return True, "Rule added."
    except Exception as e:
        logger.error(f"Error adding rule: {e}", exc_info=True)
        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))
        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: {e}", exc_info=True)
        return []

def remove_rule_entry(rule_text_to_delete: str) -> bool:
    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)
        
        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":
            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)
        return False

def get_all_rules_cached() -> list[str]:
    if not _initialized: initialize_memory_system()
    return list(_rules_items_list)

def get_all_memories_cached() -> list[dict]:
    if not _initialized: initialize_memory_system()
    valid_mems = []
    for m_str in _memory_items_list:
        try:
            valid_mems.append(json.loads(m_str))
        except json.JSONDecodeError:
            continue # Skip corrupted data for UI display
    return valid_mems

def clear_all_memory_data_backend() -> bool:
    global _memory_items_list, _faiss_memory_index
    if not _initialized: initialize_memory_system()
    _memory_items_list.clear()
    if _faiss_memory_index: _faiss_memory_index.reset()
    try:
        if STORAGE_BACKEND == "SQLITE":
            with _get_sqlite_connection() as conn: conn.execute("DELETE FROM memories"); conn.commit()
        elif STORAGE_BACKEND == "HF_DATASET":
            Dataset.from_dict({"memory_json": []}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
        return True
    except Exception as e:
        logger.error(f"Error clearing memory data: {e}"); return False

def clear_all_rules_data_backend() -> bool:
    global _rules_items_list, _faiss_rules_index
    if not _initialized: initialize_memory_system()
    _rules_items_list.clear()
    if _faiss_rules_index: _faiss_rules_index.reset()
    try:
        if STORAGE_BACKEND == "SQLITE":
            with _get_sqlite_connection() as conn: conn.execute("DELETE FROM rules"); conn.commit()
        elif STORAGE_BACKEND == "HF_DATASET":
            Dataset.from_dict({"rule_text": []}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
        return True
    except Exception as e:
        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)
    if _faiss_memory_index: faiss.write_index(_faiss_memory_index, os.path.join(faiss_dir, "memory_index.faiss"))
    if _faiss_rules_index: faiss.write_index(_faiss_rules_index, os.path.join(faiss_dir, "rules_index.faiss"))