Update model_logic.py

model_logic.py

@@ -1,437 +1,506 @@
-# model_handler.py
+# memory_logic.py
 import os
-import requests
 import json
+import time
+from datetime import datetime
 import logging
-from dotenv import load_dotenv
+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.")
 
-# Load environment variables from .env file
-load_dotenv()
 
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-)
 logger = logging.getLogger(__name__)
-
-# Maps provider name (uppercase) to environment variable name for API key
-API_KEYS_ENV_VARS = {
-  "HUGGINGFACE": 'HF_TOKEN', # Note: HF_TOKEN is often used for general HF auth
-  "GROQ": 'GROQ_API_KEY',
-  "OPENROUTER": 'OPENROUTER_API_KEY',
-  "TOGETHERAI": 'TOGETHERAI_API_KEY',
-  "COHERE": 'COHERE_API_KEY',
-  "XAI": 'XAI_API_KEY',
-  "OPENAI": 'OPENAI_API_KEY',
-  "GOOGLE": 'GOOGLE_API_KEY', # Or GOOGLE_GEMINI_API_KEY etc.
-}
-
-API_URLS = {
-  "HUGGINGFACE": 'https://api-inference.huggingface.co/models/',
-  "GROQ": 'https://api.groq.com/openai/v1/chat/completions',
-  "OPENROUTER": 'https://openrouter.ai/api/v1/chat/completions',
-  "TOGETHERAI": 'https://api.together.ai/v1/chat/completions',
-  "COHERE": 'https://api.cohere.ai/v1/chat', # v1 is common for chat, was v2 in ai-learn
-  "XAI": 'https://api.x.ai/v1/chat/completions',
-  "OPENAI": 'https://api.openai.com/v1/chat/completions',
-  "GOOGLE": 'https://generativelanguage.googleapis.com/v1beta/models/',
-}
-
-# MODELS_BY_PROVIDER = json.load(open("./models.json"))  ## commented for demo
-MODELS_BY_PROVIDER = {
-    "groq": {
-        "default": "llama3-8b-8192",
-        "models": {
-            "Llama 3 8B (Groq)": "llama3-8b-8192",
-        }
-    }
-}
-
-def _get_api_key(provider: str, ui_api_key_override: str = None) -> str | None:
-    """
-    Retrieves API key for a given provider.
-    Priority: UI Override > Environment Variable from API_KEYS_ENV_VARS > Specific (e.g. HF_TOKEN for HuggingFace).
-    """
-    provider_upper = provider.upper()
-    if ui_api_key_override and ui_api_key_override.strip():
-        logger.debug(f"Using UI-provided API key for {provider_upper}.")
-        return ui_api_key_override.strip()
-
-    env_var_name = API_KEYS_ENV_VARS.get(provider_upper)
-    if env_var_name:
-        env_key = os.getenv(env_var_name)
-        if env_key and env_key.strip():
-            logger.debug(f"Using API key from env var '{env_var_name}' for {provider_upper}.")
-            return env_key.strip()
-
-    # Specific fallback for HuggingFace if HF_TOKEN is set and API_KEYS_ENV_VARS['HUGGINGFACE'] wasn't specific enough
-    if provider_upper == 'HUGGINGFACE':
-         hf_token_fallback = os.getenv("HF_TOKEN")
-         if hf_token_fallback and hf_token_fallback.strip():
-             logger.debug("Using HF_TOKEN as fallback for HuggingFace provider.")
-             return hf_token_fallback.strip()
-
-    logger.warning(f"API Key not found for provider '{provider_upper}'. Checked UI override and environment variable '{env_var_name or 'N/A'}'.")
-    return None
-
-def get_available_providers() -> list[str]:
-    """Returns a sorted list of available provider names (e.g., 'groq', 'openai')."""
-    return sorted(list(MODELS_BY_PROVIDER.keys()))
-
-def get_model_display_names_for_provider(provider: str) -> list[str]:
-    """Returns a sorted list of model display names for a given provider."""
-    return sorted(list(MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {}).keys()))
-
-def get_default_model_display_name_for_provider(provider: str) -> str | None:
-    """Gets the default model's display name for a provider."""
-    provider_data = MODELS_BY_PROVIDER.get(provider.lower(), {})
-    models_dict = provider_data.get("models", {})
-    default_model_id = provider_data.get("default")
-
-    if default_model_id and models_dict:
-        for display_name, model_id_val in models_dict.items():
-            if model_id_val == default_model_id:
-                return display_name
-    
-    # Fallback to the first model in the sorted list if default not found or not set
-    if models_dict:
-        sorted_display_names = sorted(list(models_dict.keys()))
-        if sorted_display_names:
-            return sorted_display_names[0]
-    return None
-
-def get_model_id_from_display_name(provider: str, display_name: str) -> str | None:
-    """Gets the actual model ID from its display name for a given provider."""
-    models = MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {})
-    return models.get(display_name)
-
-
-def call_model_stream(provider: str, model_display_name: str, messages: list[dict], api_key_override: str = None, temperature: float = 0.7, max_tokens: int = None) -> iter:
-    """
-    Calls the specified model via its provider and streams the response.
-    Handles provider-specific request formatting and error handling.
-    Yields chunks of the response text or an error string.
-    """
-    provider_lower = provider.lower()
-    api_key = _get_api_key(provider_lower, api_key_override)
-    base_url = API_URLS.get(provider.upper())
-    model_id = get_model_id_from_display_name(provider_lower, model_display_name)
-
-    if not api_key:
-        env_var_name = API_KEYS_ENV_VARS.get(provider.upper(), 'N/A')
-        yield f"Error: API Key not found for {provider}. Please set it in the UI or env var '{env_var_name}'."
+# Suppress verbose logs from dependencies
+for lib_name in ["sentence_transformers", "faiss", "datasets", "huggingface_hub"]:
+    if logging.getLogger(lib_name): # Check if logger exists
+        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") # Changed default path
+HF_TOKEN = os.getenv("HF_TOKEN")
+HF_MEMORY_DATASET_REPO = os.getenv("HF_MEMORY_DATASET_REPO", "broadfield-dev/ai-brain") # Example
+HF_RULES_DATASET_REPO = os.getenv("HF_RULES_DATASET_REPO", "broadfield-dev/ai-rules")     # Example
+
+# --- Globals for RAG within this module ---
+_embedder = None
+_dimension = 384 # Default, will be set by embedder
+_faiss_memory_index = None
+_memory_items_list = []  # Stores JSON strings of memory objects for RAM, or loaded from DB/HF
+_faiss_rules_index = None
+_rules_items_list = []   # Stores rule text strings
+
+_initialized = False
+_init_lock = threading.Lock()
+
+# --- Helper: SQLite Connection ---
+def _get_sqlite_connection():
+    if not sqlite3:
+        raise ImportError("sqlite3 module is required for SQLite backend but not found.")
+    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
+
+def _init_sqlite_tables():
+    if STORAGE_BACKEND != "SQLITE" or not sqlite3:
         return
-    if not base_url:
-        yield f"Error: Unknown provider '{provider}' or missing API URL configuration."
-        return
-    if not model_id:
-         yield f"Error: Model ID not found for '{model_display_name}' under provider '{provider}'. Check configuration."
-         return
+    try:
+        with _get_sqlite_connection() as conn:
+            cursor = conn.cursor()
+            # Stores JSON string of the memory object
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS memories (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    memory_json TEXT NOT NULL,
+                    # Optionally add embedding here if not using separate FAISS index
+                    # embedding BLOB, 
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            # Stores the rule text directly
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS rules (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    rule_text TEXT NOT NULL UNIQUE,
+                    # embedding BLOB,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            conn.commit()
+        logger.info("SQLite tables for memories and rules checked/created.")
+    except Exception as e:
+        logger.error(f"SQLite table initialization error: {e}", exc_info=True)
+
+# --- 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
 
-    headers = {}
-    payload = {}
-    request_url = base_url
+        logger.info(f"Initializing memory system with backend: {STORAGE_BACKEND}")
+        init_start_time = time.time()
 
-    logger.info(f"Streaming from {provider}/{model_display_name} (ID: {model_id})...")
+        # 1. Load Sentence Transformer Model (always needed for semantic operations)
+        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 # Mark as not properly initialized
+            return
+        
+        if not _embedder:
+            try:
+                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 # Cannot proceed without embedder
+
+        # 2. Initialize SQLite if used
+        if STORAGE_BACKEND == "SQLITE":
+            _init_sqlite_tables()
+
+        # 3. Load Memories
+        logger.info("Loading memories...")
+        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")]
+            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:
+            try:
+                logger.info(f"Attempting to load 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:
+                    temp_memories_json = [m_json for m_json in dataset["train"]["memory_json"] if isinstance(m_json, str)]
+                else: logger.warning(f"HF Dataset {HF_MEMORY_DATASET_REPO} for memories not found or 'memory_json' column missing.")
+            except Exception as e: logger.error(f"Error loading memories from HF Dataset ({HF_MEMORY_DATASET_REPO}): {e}")
+        
+        _memory_items_list = temp_memories_json
+        logger.info(f"Loaded {len(_memory_items_list)} memory items from {STORAGE_BACKEND}.")
+
+        # 4. Build/Load FAISS Memory Index
+        _faiss_memory_index = faiss.IndexFlatL2(_dimension)
+        if _memory_items_list:
+            logger.info(f"Building FAISS index for {len(_memory_items_list)} memories...")
+            # Extract text to embed from memory JSON objects
+            texts_to_embed_mem = []
+            for mem_json_str in _memory_items_list:
+                try:
+                    mem_obj = json.loads(mem_json_str)
+                    # Consistent embedding strategy: user input + bot response + takeaway
+                    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'}")
+
+
+        # 5. Load Rules
+        logger.info("Loading rules...")
+        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")]
+            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:
+            try:
+                logger.info(f"Attempting to load 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:
+                    temp_rules_text = [r_text for r_text in dataset["train"]["rule_text"] if isinstance(r_text, str) and r_text.strip()]
+                else: logger.warning(f"HF Dataset {HF_RULES_DATASET_REPO} for rules not found or 'rule_text' column missing.")
+            except Exception as e: logger.error(f"Error loading rules from HF Dataset ({HF_RULES_DATASET_REPO}): {e}")
+        
+        _rules_items_list = sorted(list(set(temp_rules_text))) # Ensure unique and sorted
+        logger.info(f"Loaded {len(_rules_items_list)} rule items from {STORAGE_BACKEND}.")
+
+        # 6. Build/Load FAISS Rules Index
+        _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: # Check again in case it became empty after filtering
+                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'}")
+
+        _initialized = True
+        logger.info(f"Memory system initialization complete in {time.time() - init_start_time:.2f}s")
+
+
+# --- Memory Operations (Semantic) ---
+def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
+    """Adds a memory entry to the configured backend and FAISS index."""
+    global _memory_items_list, _faiss_memory_index
+    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."
+
+    memory_obj = {
+        "user_input": user_input,
+        "metrics": metrics,
+        "bot_response": bot_response,
+        "timestamp": datetime.utcnow().isoformat()
+    }
+    memory_json_str = json.dumps(memory_obj)
     
-    # --- Standard OpenAI-compatible providers ---
-    if provider_lower in ["groq", "openrouter", "togetherai", "openai", "xai"]:
-        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-        payload = {"model": model_id, "messages": messages, "stream": True, "temperature": temperature}
-        if max_tokens: payload["max_tokens"] = max_tokens
-
-        if provider_lower == "openrouter":
-             headers["HTTP-Referer"] = os.getenv("OPENROUTER_REFERRER") or "http://localhost/gradio" # Example Referer
-             headers["X-Title"] = os.getenv("OPENROUTER_X_TITLE") or "Gradio AI Researcher"      # Example Title
+    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)
 
-        try:
-            response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
-            response.raise_for_status()
-
-            # More robust SSE parsing
-            buffer = ""
-            for chunk in response.iter_content(chunk_size=None): # Process raw bytes
-                buffer += chunk.decode('utf-8', errors='replace')
-                while '\n\n' in buffer:
-                    event_str, buffer = buffer.split('\n\n', 1)
-                    if not event_str.strip(): continue
-
-                    content_chunk = ""
-                    for line in event_str.splitlines():
-                        if line.startswith('data: '):
-                            data_json = line[len('data: '):].strip()
-                            if data_json == '[DONE]':
-                                return # Stream finished
-                            try:
-                                data = json.loads(data_json)
-                                if data.get("choices") and len(data["choices"]) > 0:
-                                    delta = data["choices"][0].get("delta", {})
-                                    if delta and delta.get("content"):
-                                        content_chunk += delta["content"]
-                            except json.JSONDecodeError:
-                                logger.warning(f"Failed to decode JSON from stream line: {data_json}")
-                    if content_chunk:
-                        yield content_chunk
-            # Process any remaining buffer content (less common with '\n\n' delimiter)
-            if buffer.strip():
-                logger.debug(f"Remaining buffer after OpenAI-like stream: {buffer}")
-
-
-        except requests.exceptions.HTTPError as e:
-            err_msg = f"API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
-            logger.error(f"{err_msg} for {provider}/{model_id}", exc_info=False)
-            yield f"Error: {err_msg}"
-        except requests.exceptions.RequestException as e:
-            logger.error(f"API Request Error for {provider}/{model_id}: {e}", exc_info=False)
-            yield f"Error: Could not connect to {provider} ({e})"
-        except Exception as e:
-            logger.exception(f"Unexpected error during {provider} stream:")
-            yield f"Error: An unexpected error occurred: {e}"
-        return
+        if embedding_np.shape != (1, _dimension):
+            logger.error(f"Memory embedding shape error: {embedding_np.shape}. Expected (1, {_dimension})")
+            return False, "Embedding shape error."
 
-    # --- Google Gemini ---
-    elif provider_lower == "google":
-        system_instruction = None
-        filtered_messages = []
-        for msg in messages:
-            if msg["role"] == "system": system_instruction = {"parts": [{"text": msg["content"]}]}
-            else:
-                role = "model" if msg["role"] == "assistant" else msg["role"]
-                filtered_messages.append({"role": role, "parts": [{"text": msg["content"]}]})
-
-        payload = {
-             "contents": filtered_messages,
-             "safetySettings": [ # Example: more permissive settings
-                 {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
-                 {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
-                 {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
-                 {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
-             ],
-             "generationConfig": {"temperature": temperature}
-        }
-        if max_tokens: payload["generationConfig"]["maxOutputTokens"] = max_tokens
-        if system_instruction: payload["system_instruction"] = system_instruction
+        # Add to FAISS
+        _faiss_memory_index.add(embedding_np)
         
-        request_url = f"{base_url}{model_id}:streamGenerateContent?key={api_key}" # API key in query param
-        headers = {"Content-Type": "application/json"}
-
-        try:
-            response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
-            response.raise_for_status()
+        # Add to in-memory list
+        _memory_items_list.append(memory_json_str)
+        
+        # Add to persistent storage
+        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+            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:
+            # This can be slow, consider batching or async push
+            logger.info(f"Pushing {len(_memory_items_list)} memories to HF Hub: {HF_MEMORY_DATASET_REPO}")
+            Dataset.from_dict({"memory_json": list(_memory_items_list)}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True) # Ensure 'private' as needed
+        
+        logger.info(f"Added memory. RAM: {len(_memory_items_list)}, FAISS: {_faiss_memory_index.ntotal}")
+        return True, "Memory added successfully."
+    except Exception as e:
+        logger.error(f"Error adding memory entry: {e}", exc_info=True)
+        # TODO: Potential rollback logic if FAISS add succeeded but backend failed (complex)
+        return False, f"Error adding memory: {e}"
+
+def retrieve_memories_semantic(query: str, k: int = 3) -> list[dict]:
+    """Retrieves k most relevant memories using semantic search."""
+    if not _initialized: initialize_memory_system()
+    if not _embedder or not _faiss_memory_index or _faiss_memory_index.ntotal == 0:
+        logger.debug("Cannot retrieve memories: Embedder, FAISS index not ready, or 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 []
             
-            # Google's stream is a bit different, often newline-delimited JSON arrays/objects
-            buffer = ""
-            for chunk in response.iter_content(chunk_size=None):
-                buffer += chunk.decode('utf-8', errors='replace')
-                # Google might send chunks that are not complete JSON objects, or multiple objects
-                # A common pattern is [ {obj1} , {obj2} ] where chunks split mid-array or mid-object.
-                # This parsing needs to be robust. A simple split by '\n' might not always work if JSON is pretty-printed.
-                # The previous code's `json.loads(f"[{decoded_line}]")` was an attempt to handle this.
-                # For now, let's assume newline delimited for simplicity, but this is a known tricky part.
-                
-                while '\n' in buffer:
-                    line, buffer = buffer.split('\n', 1)
-                    line = line.strip()
-                    if not line: continue
-                    if line.startswith(','): line = line[1:] # Handle leading commas if splitting an array
-
-                    try:
-                        # Remove "data: " prefix if present (less common for Gemini direct API but good practice)
-                        if line.startswith('data: '): line = line[len('data: '):]
-                        
-                        # Gemini often streams an array of objects, or just one object.
-                        # Try to parse as a single object first. If fails, try as array.
-                        parsed_data = None
-                        try:
-                            parsed_data = json.loads(line)
-                        except json.JSONDecodeError:
-                            # If it's part of an array, it might be missing brackets.
-                            # This heuristic is fragile. A proper SSE parser or stateful JSON parser is better.
-                            if line.startswith('{') and line.endswith('}'): # Looks like a complete object
-                                pass # already tried json.loads
-                            # Try to wrap with [] if it seems like a list content without brackets
-                            elif line.startswith('{') or line.endswith('}'):
-                                try:
-                                    temp_parsed_list = json.loads(f"[{line}]")
-                                    if temp_parsed_list and isinstance(temp_parsed_list, list):
-                                        parsed_data = temp_parsed_list[0] # take first if it becomes a list
-                                except json.JSONDecodeError:
-                                    logger.warning(f"Google: Still can't parse line even with array wrap: {line}")
-
-                        if parsed_data:
-                            data_to_process = [parsed_data] if isinstance(parsed_data, dict) else parsed_data # Ensure list
-                            for event_data in data_to_process:
-                                if not isinstance(event_data, dict): continue
-                                if event_data.get("candidates"):
-                                    for candidate in event_data["candidates"]:
-                                        if candidate.get("content", {}).get("parts"):
-                                            for part in candidate["content"]["parts"]:
-                                                if part.get("text"):
-                                                    yield part["text"]
-                    except json.JSONDecodeError:
-                        logger.warning(f"Google: JSONDecodeError for line: {line}")
-                    except Exception as e_google_proc:
-                        logger.error(f"Google: Error processing stream data: {e_google_proc}, Line: {line}")
-
-        except requests.exceptions.HTTPError as e:
-            err_msg = f"Google API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
-            logger.error(err_msg, exc_info=False)
-            yield f"Error: {err_msg}"
-        except Exception as e:
-            logger.exception(f"Unexpected error during Google stream:")
-            yield f"Error: An unexpected error occurred with Google API: {e}"
-        return
-
-    # --- Cohere ---
-    elif provider_lower == "cohere":
-        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "Accept": "application/json"}
+        distances, indices = _faiss_memory_index.search(query_embedding_np, min(k, _faiss_memory_index.ntotal))
         
-        # Cohere message format
-        chat_history_cohere = []
-        preamble_cohere = None
-        user_message_cohere = ""
-
-        temp_messages = list(messages) # Work with a copy
-        if temp_messages and temp_messages[0]["role"] == "system":
-            preamble_cohere = temp_messages.pop(0)["content"]
+        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]}...'")
+        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]:
+    """Adds a rule if valid and not a duplicate. Updates backend and FAISS."""
+    global _rules_items_list, _faiss_rules_index
+    if not _initialized: initialize_memory_system()
+    if not _embedder or not _faiss_rules_index:
+        return False, "Rule system or embedder not initialized."
+
+    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):
+        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."
+
+        _faiss_rules_index.add(embedding_np)
+        _rules_items_list.append(rule_text)
+        _rules_items_list.sort()
+
+        if STORAGE_BACKEND == "SQLITE" and sqlite3:
+            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}")
+            Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
         
-        if temp_messages:
-            user_message_cohere = temp_messages.pop()["content"] # Last message is the current user query
-            for msg in temp_messages: # Remaining are history
-                role = "USER" if msg["role"] == "user" else "CHATBOT"
-                chat_history_cohere.append({"role": role, "message": msg["content"]})
+        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)
+        # Basic rollback if FAISS add succeeded
+        if rule_text in _rules_items_list and _faiss_rules_index.ntotal > 0: # Crude check
+             # A full rollback would involve rebuilding FAISS index from _rules_items_list before append.
+             # For simplicity, this is omitted here. State could be inconsistent on error.
+             pass
+        return False, f"Error adding rule: {e}"
+
+def retrieve_rules_semantic(query: str, k: int = 5) -> list[str]:
+    """Retrieves k most relevant rules using semantic search."""
+    if not _initialized: initialize_memory_system()
+    if not _embedder or 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 not user_message_cohere:
-            yield "Error: User message is empty for Cohere."
-            return
-
-        payload = {
-            "model": model_id, 
-            "message": user_message_cohere, 
-            "stream": True, 
-            "temperature": temperature
-        }
-        if max_tokens: payload["max_tokens"] = max_tokens # Cohere uses max_tokens
-        if chat_history_cohere: payload["chat_history"] = chat_history_cohere
-        if preamble_cohere: payload["preamble"] = preamble_cohere
+        if query_embedding_np.shape[1] != _dimension: return []
+
+        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
+    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:
+    """Removes a rule from backend and rebuilds FAISS for rules."""
+    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 # Not found
+
+    try:
+        _rules_items_list.remove(rule_text_to_delete)
+        _rules_items_list.sort() # Maintain sorted order
+
+        # Rebuild FAISS index for rules (simplest way to ensure consistency after removal)
+        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: # Should not happen if list is consistent
+                logger.error("Error rebuilding FAISS for rules after removal: Embedding shape error. State might be inconsistent.")
+                # Attempt to revert _rules_items_list (add back the rule)
+                _rules_items_list.append(rule_text_to_delete)
+                _rules_items_list.sort()
+                return False # Indicate failure
+        _faiss_rules_index = new_faiss_rules_index
         
-        try:
-            response = requests.post(base_url, headers=headers, json=payload, stream=True, timeout=180)
-            response.raise_for_status()
-            
-            # Cohere SSE format is event: type\ndata: {json}\n\n
-            buffer = ""
-            for chunk_bytes in response.iter_content(chunk_size=None):
-                buffer += chunk_bytes.decode('utf-8', errors='replace')
-                while '\n\n' in buffer:
-                    event_str, buffer = buffer.split('\n\n', 1)
-                    if not event_str.strip(): continue
-                    
-                    event_type = None
-                    data_json_str = None
-                    for line in event_str.splitlines():
-                        if line.startswith("event:"): event_type = line[len("event:"):].strip()
-                        elif line.startswith("data:"): data_json_str = line[len("data:"):].strip()
-                    
-                    if data_json_str:
-                        try:
-                            data = json.loads(data_json_str)
-                            if event_type == "text-generation" and "text" in data:
-                                yield data["text"]
-                            elif event_type == "stream-end":
-                                logger.debug(f"Cohere stream ended. Finish reason: {data.get('finish_reason')}")
-                                return 
-                        except json.JSONDecodeError:
-                            logger.warning(f"Cohere: Failed to decode JSON: {data_json_str}")
-            if buffer.strip():
-                 logger.debug(f"Cohere: Remaining buffer: {buffer.strip()}")
-
-
-        except requests.exceptions.HTTPError as e:
-            err_msg = f"Cohere API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
-            logger.error(err_msg, exc_info=False)
-            yield f"Error: {err_msg}"
-        except Exception as e:
-            logger.exception(f"Unexpected error during Cohere stream:")
-            yield f"Error: An unexpected error occurred with Cohere API: {e}"
-        return
-
-    # --- HuggingFace Inference API (Basic TGI support) ---
-    # This is very basic and might not work for all models or complex scenarios.
-    # Assumes model is deployed with Text Generation Inference (TGI) and supports streaming.
-    elif provider_lower == "huggingface":
-        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-        # Construct prompt string for TGI (often needs specific formatting)
-        # This is a generic attempt, specific models might need <|user|>, <|assistant|> etc.
-        prompt_parts = []
-        for msg in messages:
-            role_prefix = ""
-            if msg['role'] == 'system': role_prefix = "System: " # Or might be ignored/handled differently
-            elif msg['role'] == 'user': role_prefix = "User: "
-            elif msg['role'] == 'assistant': role_prefix = "Assistant: "
-            prompt_parts.append(f"{role_prefix}{msg['content']}")
+        # Remove from persistent storage
+        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)
+        
+        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)
+        # Potential partial failure, state might be inconsistent.
+        return False
+
+# --- Utility functions to get all data (for UI display, etc.) ---
+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()
+    # Convert JSON strings to dicts for easier use by UI
+    mem_dicts = []
+    for mem_json_str in _memory_items_list:
+        try: mem_dicts.append(json.loads(mem_json_str))
+        except: pass # Ignore parse errors for display
+    return mem_dicts
+
+def clear_all_memory_data_backend() -> bool:
+    """Clears all memories from backend and resets in-memory FAISS/list."""
+    global _memory_items_list, _faiss_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:
+            # Deleting from HF usually means pushing an empty dataset
+            Dataset.from_dict({"memory_json": []}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
         
-        # TGI typically expects a final "Assistant: " to start generating from
-        tgi_prompt = "\n".join(prompt_parts) + "\nAssistant: "
-
-        payload = {
-            "inputs": tgi_prompt,
-            "parameters": {
-                "temperature": temperature if temperature > 0 else 0.01, # TGI needs temp > 0 for sampling
-                "max_new_tokens": max_tokens or 1024, # Default TGI max_new_tokens
-                "return_full_text": False, # We only want generated part
-                "do_sample": True if temperature > 0 else False,
-            },
-            "stream": True
-        }
-        request_url = f"{base_url}{model_id}" # Model ID is part of URL path for HF
+        _memory_items_list = []
+        if _faiss_memory_index: _faiss_memory_index.reset() # Clear FAISS index
+        logger.info("All memories cleared from backend and in-memory stores.")
+    except Exception as e:
+        logger.error(f"Error clearing all memory data: {e}")
+        success = False
+    return success
+
+def clear_all_rules_data_backend() -> bool:
+    """Clears all rules from backend and resets in-memory FAISS/list."""
+    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()
+        logger.info("All rules cleared from backend and in-memory stores.")
+    except Exception as e:
+        logger.error(f"Error clearing all rules data: {e}")
+        success = False
+    return success
+
+# Optional: Function to save FAISS indices to disk (from ai-learn, if needed for persistence between app runs with RAM backend)
+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")
+
+def save_faiss_indices_to_disk():
+    if not _initialized or not faiss: return
+    
+    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:
-            response = requests.post(request_url, headers=headers, json=payload, stream=True, timeout=180)
-            response.raise_for_status()
-
-            # TGI SSE stream: data: {"token": {"id": ..., "text": "...", "logprob": ..., "special": ...}}
-            # Or sometimes just data: "text_chunk" for simpler models/configs
-            buffer = ""
-            for chunk_bytes in response.iter_content(chunk_size=None):
-                buffer += chunk_bytes.decode('utf-8', errors='replace')
-                while '\n' in buffer: # TGI often uses single newline
-                    line, buffer = buffer.split('\n', 1)
-                    line = line.strip()
-                    if not line: continue
-
-                    if line.startswith('data:'):
-                        data_json_str = line[len('data:'):].strip()
-                        try:
-                            data = json.loads(data_json_str)
-                            if "token" in data and "text" in data["token"]:
-                                yield data["token"]["text"]
-                            elif "generated_text" in data and data.get("details") is None: # Sometimes a final non-streaming like object might appear
-                                 # This case is tricky, if it's the *only* thing then it's not really streaming
-                                 pass # For now, ignore if it's not a token object
-                            # Some TGI might send raw text if not fully SSE compliant for stream
-                            # elif isinstance(data, str): yield data 
-
-                        except json.JSONDecodeError:
-                            # If it's not JSON, it might be a raw string (less common for TGI stream=True)
-                            # For safety, only yield if it's a clear text string
-                            if not data_json_str.startswith('{') and not data_json_str.startswith('['):
-                                yield data_json_str
-                            else:
-                                logger.warning(f"HF: Failed to decode JSON and not raw string: {data_json_str}")
-            if buffer.strip():
-                 logger.debug(f"HF: Remaining buffer: {buffer.strip()}")
-
-
-        except requests.exceptions.HTTPError as e:
-            err_msg = f"HF API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
-            logger.error(err_msg, exc_info=False)
-            yield f"Error: {err_msg}"
-        except Exception as e:
-            logger.exception(f"Unexpected error during HF stream:")
-            yield f"Error: An unexpected error occurred with HF API: {e}"
-        return
+            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}")
+    
+    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}")
+
+def load_faiss_indices_from_disk():
+    global _faiss_memory_index, _faiss_rules_index
+    if not _initialized or not faiss: return
 
-    else:
-        yield f"Error: Provider '{provider}' is not configured for streaming in this handler."
-        return
+    if os.path.exists(FAISS_MEMORY_PATH) and _faiss_memory_index: # Check if index object exists
+        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).")
+            # Consistency check: FAISS ntotal vs len(_memory_items_list)
+            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.")