Update model_logic.py

model_logic.py

@@ -1,322 +1,84 @@
+# model_logic.py
 import os
+import requests
 import json
-import time
-from datetime import datetime
 import logging
-import re
-import threading
+from dotenv import load_dotenv
 
-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_dotenv()
 
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
-for lib_name in ["sentence_transformers", "faiss", "datasets", "huggingface_hub"]:
-    if logging.getLogger(lib_name): logging.getLogger(lib_name).setLevel(logging.WARNING)
-
-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")
-
-_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.")
-            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
-
-        if STORAGE_BACKEND == "SQLITE": _init_sqlite_tables()
-
-        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)
-        
-        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
-        logger.info(f"Loaded and indexed {len(_memory_items_list)} memories.")
-
-        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)
-
-        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
-        logger.info(f"Loaded and indexed {len(_rules_items_list)} rules.")
-        
-        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):
-    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()
-    if not _embedder or _faiss_memory_index is None: return False, "Memory system not ready."
+API_KEYS_ENV_VARS = {"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', "HUGGINGFACE": 'HF_TOKEN'}
+API_URLS = {"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', "XAI": 'https://api.x.ai/v1/chat/completions', "OPENAI": 'https://api.openai.com/v1/chat/completions', "GOOGLE": 'https://generativelanguage.googleapis.com/v1beta/models/', "HUGGINGFACE": 'https://api-inference.huggingface.co/models/'}
 
-    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]:
-    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."
-
-    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:
-    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
-        
-        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:
-            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
-    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
+try:
+    with open("models.json", "r") as f: MODELS_BY_PROVIDER = json.load(f)
+except (FileNotFoundError, json.JSONDecodeError):
+    logger.warning("models.json not found or invalid. Using a fallback model list.")
+    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:
+    provider_upper = provider.upper()
+    if ui_api_key_override and ui_api_key_override.strip(): 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(): return env_key.strip()
+    return None
+
+def get_available_providers() -> list[str]:
+    return sorted(list(MODELS_BY_PROVIDER.keys()))
+
+def get_model_display_names_for_provider(provider: str) -> list[str]:
+    return sorted(list(MODELS_BY_PROVIDER.get(provider.lower(), {}).get("models", {}).keys()))
+
+def get_default_model_display_name_for_provider(provider: str) -> str | None:
+    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
+    if models_dict: return sorted(list(models_dict.keys()))[0]
+    return None
+
+def get_model_id_from_display_name(provider: str, display_name: str) -> str | None:
+    return MODELS_BY_PROVIDER.get(provider.lower(), {}).get("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:
+    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 all([api_key, base_url, model_id]):
+        yield f"Error: Configuration missing for {provider}/{model_display_name}."
+        return
+
+    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", "http://localhost")
 
-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
+        response = requests.post(base_url, headers=headers, json=payload, stream=True, timeout=180)
+        response.raise_for_status()
+        buffer = ""
+        for chunk in response.iter_content(chunk_size=None):
+            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() or not event_str.startswith('data: '): continue
+                data_json = event_str[len('data: '):].strip()
+                if data_json == '[DONE]': return
+                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"): yield delta["content"]
+                except json.JSONDecodeError: continue
+    except requests.exceptions.HTTPError as e:
+        yield f"Error: API HTTP Error ({e.response.status_code}): {e.response.text[:200]}"
     except Exception as e:
-        logger.error(f"Error clearing rules data: {e}"); return False
-
-def save_faiss_indices_to_disk():
-    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"))
+        yield f"Error: An unexpected error occurred: {e}"