broadfield-dev commited on
Commit
03a2414
·
verified ·
1 Parent(s): 727e0fe

Update memory_logic.py

Browse files
Files changed (1) hide show
  1. memory_logic.py +468 -123
memory_logic.py CHANGED
@@ -1,163 +1,508 @@
 
1
  import os
2
  import json
3
- import logging
4
  from datetime import datetime
5
- import re # For insight format validation
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
  logger = logging.getLogger(__name__)
 
 
 
 
 
8
 
9
- DATA_DIR = "app_data"
10
- MEMORIES_FILE = os.path.join(DATA_DIR, "conversation_memories.jsonl") # JSON Lines format
11
- RULES_FILE = os.path.join(DATA_DIR, "learned_rules.jsonl") # Rules/Insights, also JSON Lines
 
 
 
12
 
13
- # Ensure data directory exists
14
- os.makedirs(DATA_DIR, exist_ok=True)
 
 
 
 
 
15
 
16
- # --- Rules/Insights Management ---
 
17
 
18
- def load_rules_from_file() -> list[str]:
19
- """Loads rules (insights) from the JSON Lines file."""
20
- rules = []
21
- if not os.path.exists(RULES_FILE):
22
- return rules
 
 
 
 
 
 
 
23
  try:
24
- with open(RULES_FILE, 'r', encoding='utf-8') as f:
25
- for line in f:
26
- if line.strip():
27
- try:
28
- # Assuming each line is a JSON object like {"rule_text": "...", "timestamp": "..."}
29
- # For simplicity, if we only stored the text previously, adapt here.
30
- # Let's assume we store {"text": "rule_text_content"}
31
- data = json.loads(line)
32
- if "text" in data and isinstance(data["text"], str) and data["text"].strip():
33
- rules.append(data["text"].strip())
34
- elif isinstance(data, str): # If old format was just text per line
35
- rules.append(data.strip())
36
-
37
- except json.JSONDecodeError:
38
- logger.warning(f"Skipping malformed JSON line in rules file: {line.strip()}")
39
- logger.info(f"Loaded {len(rules)} rules from {RULES_FILE}")
 
 
 
 
 
 
 
40
  except Exception as e:
41
- logger.error(f"Error loading rules from {RULES_FILE}: {e}", exc_info=True)
42
- return sorted(list(set(rules))) # Ensure unique and sorted
43
 
44
- def save_rule_to_file(rule_text: str) -> bool:
45
- """Saves a single rule (insight) to the JSON Lines file if it's new and valid."""
46
- rule_text = rule_text.strip()
47
- if not rule_text:
48
- logger.warning("Attempted to save an empty rule.")
49
- return False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
 
51
- # Validate format: [TYPE|SCORE] Text
52
- if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\](.*)", rule_text, re.I|re.DOTALL):
53
- logger.warning(f"Rule '{rule_text[:50]}...' has invalid format. Not saving.")
54
- return False
55
 
56
- current_rules = load_rules_from_file()
57
- if rule_text in current_rules:
58
- logger.info(f"Rule '{rule_text[:50]}...' already exists. Not saving duplicate.")
59
- return False # Or True if "already exists" is considered success
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61
  try:
62
- with open(RULES_FILE, 'a', encoding='utf-8') as f:
63
- # Store as JSON object for potential future metadata
64
- json.dump({"text": rule_text, "added_at": datetime.utcnow().isoformat()}, f)
65
- f.write('\n')
66
- logger.info(f"Saved new rule: {rule_text[:70]}...")
67
- return True
68
- except Exception as e:
69
- logger.error(f"Error saving rule '{rule_text[:50]}...' to {RULES_FILE}: {e}", exc_info=True)
70
- return False
71
 
72
- def delete_rule_from_file(rule_text_to_delete: str) -> bool:
73
- """Deletes a rule from the file."""
74
- rule_text_to_delete = rule_text_to_delete.strip()
75
- if not rule_text_to_delete: return False
76
 
77
- current_rules = load_rules_from_file()
78
- if rule_text_to_delete not in current_rules:
79
- logger.info(f"Rule '{rule_text_to_delete[:50]}...' not found for deletion.")
80
- return False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
81
 
82
- updated_rules = [rule for rule in current_rules if rule != rule_text_to_delete]
83
  try:
84
- with open(RULES_FILE, 'w', encoding='utf-8') as f: # Overwrite with updated list
85
- for rule_text in updated_rules:
86
- json.dump({"text": rule_text, "added_at": "unknown"}, f) # timestamp lost on rewrite this way
87
- f.write('\n')
88
- logger.info(f"Deleted rule: {rule_text_to_delete[:70]}...")
89
- return True
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
90
  except Exception as e:
91
- logger.error(f"Error deleting rule '{rule_text_to_delete[:50]}...' from {RULES_FILE}: {e}", exc_info=True)
92
- return False
93
 
94
 
95
- # --- Conversation Memories Management ---
 
 
 
 
 
 
 
 
 
 
 
 
 
96
 
97
- def load_memories_from_file() -> list[dict]:
98
- """Loads conversation memories from the JSON Lines file."""
99
- memories = []
100
- if not os.path.exists(MEMORIES_FILE):
101
- return memories
102
  try:
103
- with open(MEMORIES_FILE, 'r', encoding='utf-8') as f:
104
- for line in f:
105
- if line.strip():
106
- try:
107
- mem_obj = json.loads(line)
108
- # Basic validation for expected keys
109
- if all(k in mem_obj for k in ["user_input", "bot_response", "metrics", "timestamp"]):
110
- memories.append(mem_obj)
111
- else:
112
- logger.warning(f"Skipping memory object with missing keys: {line.strip()}")
113
- except json.JSONDecodeError:
114
- logger.warning(f"Skipping malformed JSON line in memories file: {line.strip()}")
115
- logger.info(f"Loaded {len(memories)} memories from {MEMORIES_FILE}")
 
 
 
 
 
 
 
116
  except Exception as e:
117
- logger.error(f"Error loading memories from {MEMORIES_FILE}: {e}", exc_info=True)
118
- # Sort by timestamp if needed, though append-only usually keeps order
119
- return sorted(memories, key=lambda x: x.get("timestamp", ""))
 
 
 
 
120
 
 
 
 
 
 
 
 
 
 
 
121
 
122
- def save_memory_to_file(user_input: str, bot_response: str, metrics: dict) -> bool:
123
- """Saves a conversation memory to the JSON Lines file."""
124
- if not user_input or not bot_response: # Metrics can be empty
125
- logger.warning("Attempted to save memory with empty user input or bot response.")
126
- return False
 
 
127
 
128
- memory_entry = {
129
- "user_input": user_input,
130
- "bot_response": bot_response,
131
- "metrics": metrics,
132
- "timestamp": datetime.utcnow().isoformat()
133
- }
 
 
 
134
 
135
  try:
136
- with open(MEMORIES_FILE, 'a', encoding='utf-8') as f:
137
- json.dump(memory_entry, f)
138
- f.write('\n')
139
- logger.info(f"Saved new memory. User: {user_input[:50]}...")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
140
  return True
141
  except Exception as e:
142
- logger.error(f"Error saving memory to {MEMORIES_FILE}: {e}", exc_info=True)
 
143
  return False
144
 
145
- def clear_all_rules() -> bool:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  try:
147
- if os.path.exists(RULES_FILE):
148
- os.remove(RULES_FILE)
149
- logger.info("All rules cleared.")
150
- return True
 
 
 
 
 
151
  except Exception as e:
152
- logger.error(f"Error clearing rules file {RULES_FILE}: {e}")
153
- return False
 
154
 
155
- def clear_all_memories() -> bool:
 
 
 
 
 
156
  try:
157
- if os.path.exists(MEMORIES_FILE):
158
- os.remove(MEMORIES_FILE)
159
- logger.info("All memories cleared.")
160
- return True
 
 
 
 
161
  except Exception as e:
162
- logger.error(f"Error clearing memories file {MEMORIES_FILE}: {e}")
163
- return False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # memory_logic.py
2
  import os
3
  import json
4
+ import time
5
  from datetime import datetime
6
+ import logging
7
+ import re
8
+ import threading
9
+
10
+ # Conditionally import heavy dependencies
11
+ try:
12
+ from sentence_transformers import SentenceTransformer
13
+ import faiss
14
+ import numpy as np
15
+ except ImportError:
16
+ SentenceTransformer, faiss, np = None, None, None
17
+ logging.warning("SentenceTransformers, FAISS, or NumPy not installed. Semantic search will be unavailable.")
18
+
19
+ try:
20
+ import sqlite3
21
+ except ImportError:
22
+ sqlite3 = None
23
+ logging.warning("sqlite3 module not available. SQLite backend will be unavailable.")
24
+
25
+ try:
26
+ from datasets import load_dataset, Dataset
27
+ except ImportError:
28
+ load_dataset, Dataset = None, None
29
+ logging.warning("datasets library not installed. Hugging Face Dataset backend will be unavailable.")
30
+
31
 
32
  logger = logging.getLogger(__name__)
33
+ # Suppress verbose logs from dependencies
34
+ for lib_name in ["sentence_transformers", "faiss", "datasets", "huggingface_hub"]:
35
+ if logging.getLogger(lib_name): # Check if logger exists
36
+ logging.getLogger(lib_name).setLevel(logging.WARNING)
37
+
38
 
39
+ # --- Configuration (Read directly from environment variables) ---
40
+ STORAGE_BACKEND = os.getenv("STORAGE_BACKEND", "RAM").upper()
41
+ SQLITE_DB_PATH = os.getenv("SQLITE_DB_PATH", "app_data/ai_memory.db") # Changed default path
42
+ HF_TOKEN = os.getenv("HF_TOKEN")
43
+ HF_MEMORY_DATASET_REPO = os.getenv("HF_MEMORY_DATASET_REPO", "your_hf_username/app-memories") # Example
44
+ HF_RULES_DATASET_REPO = os.getenv("HF_RULES_DATASET_REPO", "your_hf_username/app-rules") # Example
45
 
46
+ # --- Globals for RAG within this module ---
47
+ _embedder = None
48
+ _dimension = 384 # Default, will be set by embedder
49
+ _faiss_memory_index = None
50
+ _memory_items_list = [] # Stores JSON strings of memory objects for RAM, or loaded from DB/HF
51
+ _faiss_rules_index = None
52
+ _rules_items_list = [] # Stores rule text strings
53
 
54
+ _initialized = False
55
+ _init_lock = threading.Lock()
56
 
57
+ # --- Helper: SQLite Connection ---
58
+ def _get_sqlite_connection():
59
+ if not sqlite3:
60
+ raise ImportError("sqlite3 module is required for SQLite backend but not found.")
61
+ db_dir = os.path.dirname(SQLITE_DB_PATH)
62
+ if db_dir and not os.path.exists(db_dir):
63
+ os.makedirs(db_dir, exist_ok=True)
64
+ return sqlite3.connect(SQLITE_DB_PATH, timeout=10) # Added timeout
65
+
66
+ def _init_sqlite_tables():
67
+ if STORAGE_BACKEND != "SQLITE" or not sqlite3:
68
+ return
69
  try:
70
+ with _get_sqlite_connection() as conn:
71
+ cursor = conn.cursor()
72
+ # Stores JSON string of the memory object
73
+ cursor.execute("""
74
+ CREATE TABLE IF NOT EXISTS memories (
75
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
76
+ memory_json TEXT NOT NULL,
77
+ # Optionally add embedding here if not using separate FAISS index
78
+ # embedding BLOB,
79
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
80
+ )
81
+ """)
82
+ # Stores the rule text directly
83
+ cursor.execute("""
84
+ CREATE TABLE IF NOT EXISTS rules (
85
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
86
+ rule_text TEXT NOT NULL UNIQUE,
87
+ # embedding BLOB,
88
+ created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
89
+ )
90
+ """)
91
+ conn.commit()
92
+ logger.info("SQLite tables for memories and rules checked/created.")
93
  except Exception as e:
94
+ logger.error(f"SQLite table initialization error: {e}", exc_info=True)
 
95
 
96
+ # --- Initialization ---
97
+ def initialize_memory_system():
98
+ global _initialized, _embedder, _dimension, _faiss_memory_index, _memory_items_list, _faiss_rules_index, _rules_items_list
99
+
100
+ with _init_lock:
101
+ if _initialized:
102
+ logger.info("Memory system already initialized.")
103
+ return
104
+
105
+ logger.info(f"Initializing memory system with backend: {STORAGE_BACKEND}")
106
+ init_start_time = time.time()
107
+
108
+ # 1. Load Sentence Transformer Model (always needed for semantic operations)
109
+ if not SentenceTransformer or not faiss or not np:
110
+ logger.error("Core RAG libraries (SentenceTransformers, FAISS, NumPy) not available. Cannot initialize semantic memory.")
111
+ _initialized = False # Mark as not properly initialized
112
+ return
113
+
114
+ if not _embedder:
115
+ try:
116
+ logger.info("Loading SentenceTransformer model (all-MiniLM-L6-v2)...")
117
+ _embedder = SentenceTransformer('all-MiniLM-L6-v2', cache_folder="./sentence_transformer_cache")
118
+ _dimension = _embedder.get_sentence_embedding_dimension() or 384
119
+ logger.info(f"SentenceTransformer loaded. Dimension: {_dimension}")
120
+ except Exception as e:
121
+ logger.critical(f"FATAL: Error loading SentenceTransformer: {e}", exc_info=True)
122
+ _initialized = False
123
+ return # Cannot proceed without embedder
124
+
125
+ # 2. Initialize SQLite if used
126
+ if STORAGE_BACKEND == "SQLITE":
127
+ _init_sqlite_tables()
128
+
129
+ # 3. Load Memories
130
+ logger.info("Loading memories...")
131
+ temp_memories_json = []
132
+ if STORAGE_BACKEND == "RAM":
133
+ _memory_items_list = [] # Start fresh for RAM backend
134
+ elif STORAGE_BACKEND == "SQLITE" and sqlite3:
135
+ try:
136
+ with _get_sqlite_connection() as conn:
137
+ temp_memories_json = [row[0] for row in conn.execute("SELECT memory_json FROM memories ORDER BY created_at ASC")]
138
+ except Exception as e: logger.error(f"Error loading memories from SQLite: {e}")
139
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset and load_dataset:
140
+ try:
141
+ logger.info(f"Attempting to load memories from HF Dataset: {HF_MEMORY_DATASET_REPO}")
142
+ dataset = load_dataset(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True) # Add download_mode if needed
143
+ if "train" in dataset and "memory_json" in dataset["train"].column_names: # Assuming 'memory_json' column
144
+ temp_memories_json = [m_json for m_json in dataset["train"]["memory_json"] if isinstance(m_json, str)]
145
+ else: logger.warning(f"HF Dataset {HF_MEMORY_DATASET_REPO} for memories not found or 'memory_json' column missing.")
146
+ except Exception as e: logger.error(f"Error loading memories from HF Dataset ({HF_MEMORY_DATASET_REPO}): {e}")
147
+
148
+ _memory_items_list = temp_memories_json
149
+ logger.info(f"Loaded {len(_memory_items_list)} memory items from {STORAGE_BACKEND}.")
150
+
151
+ # 4. Build/Load FAISS Memory Index
152
+ _faiss_memory_index = faiss.IndexFlatL2(_dimension)
153
+ if _memory_items_list:
154
+ logger.info(f"Building FAISS index for {len(_memory_items_list)} memories...")
155
+ # Extract text to embed from memory JSON objects
156
+ texts_to_embed_mem = []
157
+ for mem_json_str in _memory_items_list:
158
+ try:
159
+ mem_obj = json.loads(mem_json_str)
160
+ # Consistent embedding strategy: user input + bot response + takeaway
161
+ text = f"User: {mem_obj.get('user_input','')}\nAI: {mem_obj.get('bot_response','')}\nTakeaway: {mem_obj.get('metrics',{}).get('takeaway','N/A')}"
162
+ texts_to_embed_mem.append(text)
163
+ except json.JSONDecodeError:
164
+ logger.warning(f"Skipping malformed memory JSON for FAISS indexing: {mem_json_str[:100]}")
165
+
166
+ if texts_to_embed_mem:
167
+ try:
168
+ embeddings = _embedder.encode(texts_to_embed_mem, convert_to_tensor=False, show_progress_bar=False) # convert_to_numpy=True
169
+ embeddings_np = np.array(embeddings, dtype=np.float32)
170
+ if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(texts_to_embed_mem) and embeddings_np.shape[1] == _dimension:
171
+ _faiss_memory_index.add(embeddings_np)
172
+ 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'}")
173
+ except Exception as e_faiss_mem: logger.error(f"Error building FAISS memory index: {e_faiss_mem}")
174
+ logger.info(f"FAISS memory index built. Total items: {_faiss_memory_index.ntotal if _faiss_memory_index else 'N/A'}")
175
 
 
 
 
 
176
 
177
+ # 5. Load Rules
178
+ logger.info("Loading rules...")
179
+ temp_rules_text = []
180
+ if STORAGE_BACKEND == "RAM":
181
+ _rules_items_list = []
182
+ elif STORAGE_BACKEND == "SQLITE" and sqlite3:
183
+ try:
184
+ with _get_sqlite_connection() as conn:
185
+ temp_rules_text = [row[0] for row in conn.execute("SELECT rule_text FROM rules ORDER BY created_at ASC")]
186
+ except Exception as e: logger.error(f"Error loading rules from SQLite: {e}")
187
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset and load_dataset:
188
+ try:
189
+ logger.info(f"Attempting to load rules from HF Dataset: {HF_RULES_DATASET_REPO}")
190
+ dataset = load_dataset(HF_RULES_DATASET_REPO, token=HF_TOKEN, trust_remote_code=True)
191
+ if "train" in dataset and "rule_text" in dataset["train"].column_names:
192
+ temp_rules_text = [r_text for r_text in dataset["train"]["rule_text"] if isinstance(r_text, str) and r_text.strip()]
193
+ else: logger.warning(f"HF Dataset {HF_RULES_DATASET_REPO} for rules not found or 'rule_text' column missing.")
194
+ except Exception as e: logger.error(f"Error loading rules from HF Dataset ({HF_RULES_DATASET_REPO}): {e}")
195
+
196
+ _rules_items_list = sorted(list(set(temp_rules_text))) # Ensure unique and sorted
197
+ logger.info(f"Loaded {len(_rules_items_list)} rule items from {STORAGE_BACKEND}.")
198
 
199
+ # 6. Build/Load FAISS Rules Index
200
+ _faiss_rules_index = faiss.IndexFlatL2(_dimension)
201
+ if _rules_items_list:
202
+ logger.info(f"Building FAISS index for {len(_rules_items_list)} rules...")
203
+ if _rules_items_list: # Check again in case it became empty after filtering
204
+ try:
205
+ embeddings = _embedder.encode(_rules_items_list, convert_to_tensor=False, show_progress_bar=False)
206
+ embeddings_np = np.array(embeddings, dtype=np.float32)
207
+ if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(_rules_items_list) and embeddings_np.shape[1] == _dimension:
208
+ _faiss_rules_index.add(embeddings_np)
209
+ 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'}")
210
+ except Exception as e_faiss_rule: logger.error(f"Error building FAISS rule index: {e_faiss_rule}")
211
+ logger.info(f"FAISS rules index built. Total items: {_faiss_rules_index.ntotal if _faiss_rules_index else 'N/A'}")
212
+
213
+ _initialized = True
214
+ logger.info(f"Memory system initialization complete in {time.time() - init_start_time:.2f}s")
215
+
216
+
217
+ # --- Memory Operations (Semantic) ---
218
+ def add_memory_entry(user_input: str, metrics: dict, bot_response: str) -> tuple[bool, str]:
219
+ """Adds a memory entry to the configured backend and FAISS index."""
220
+ global _memory_items_list, _faiss_memory_index
221
+ if not _initialized: initialize_memory_system()
222
+ if not _embedder or not _faiss_memory_index:
223
+ return False, "Memory system or embedder not initialized for adding memory."
224
+
225
+ memory_obj = {
226
+ "user_input": user_input,
227
+ "metrics": metrics,
228
+ "bot_response": bot_response,
229
+ "timestamp": datetime.utcnow().isoformat()
230
+ }
231
+ memory_json_str = json.dumps(memory_obj)
232
+
233
+ text_to_embed = f"User: {user_input}\nAI: {bot_response}\nTakeaway: {metrics.get('takeaway', 'N/A')}"
234
+
235
  try:
236
+ embedding = _embedder.encode([text_to_embed], convert_to_tensor=False)
237
+ embedding_np = np.array(embedding, dtype=np.float32).reshape(1, -1)
 
 
 
 
 
 
 
238
 
239
+ if embedding_np.shape != (1, _dimension):
240
+ logger.error(f"Memory embedding shape error: {embedding_np.shape}. Expected (1, {_dimension})")
241
+ return False, "Embedding shape error."
 
242
 
243
+ # Add to FAISS
244
+ _faiss_memory_index.add(embedding_np)
245
+
246
+ # Add to in-memory list
247
+ _memory_items_list.append(memory_json_str)
248
+
249
+ # Add to persistent storage
250
+ if STORAGE_BACKEND == "SQLITE" and sqlite3:
251
+ with _get_sqlite_connection() as conn:
252
+ conn.execute("INSERT INTO memories (memory_json) VALUES (?)", (memory_json_str,))
253
+ conn.commit()
254
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
255
+ # This can be slow, consider batching or async push
256
+ logger.info(f"Pushing {len(_memory_items_list)} memories to HF Hub: {HF_MEMORY_DATASET_REPO}")
257
+ 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
258
+
259
+ logger.info(f"Added memory. RAM: {len(_memory_items_list)}, FAISS: {_faiss_memory_index.ntotal}")
260
+ return True, "Memory added successfully."
261
+ except Exception as e:
262
+ logger.error(f"Error adding memory entry: {e}", exc_info=True)
263
+ # TODO: Potential rollback logic if FAISS add succeeded but backend failed (complex)
264
+ return False, f"Error adding memory: {e}"
265
+
266
+ def retrieve_memories_semantic(query: str, k: int = 3) -> list[dict]:
267
+ """Retrieves k most relevant memories using semantic search."""
268
+ if not _initialized: initialize_memory_system()
269
+ if not _embedder or not _faiss_memory_index or _faiss_memory_index.ntotal == 0:
270
+ logger.debug("Cannot retrieve memories: Embedder, FAISS index not ready, or index is empty.")
271
+ return []
272
 
 
273
  try:
274
+ query_embedding = _embedder.encode([query], convert_to_tensor=False)
275
+ query_embedding_np = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
276
+
277
+ if query_embedding_np.shape[1] != _dimension:
278
+ logger.error(f"Query embedding dimension mismatch. Expected {_dimension}, got {query_embedding_np.shape[1]}")
279
+ return []
280
+
281
+ distances, indices = _faiss_memory_index.search(query_embedding_np, min(k, _faiss_memory_index.ntotal))
282
+
283
+ results = []
284
+ for i in indices[0]:
285
+ if 0 <= i < len(_memory_items_list):
286
+ try:
287
+ results.append(json.loads(_memory_items_list[i]))
288
+ except json.JSONDecodeError:
289
+ logger.warning(f"Could not parse memory JSON from list at index {i}")
290
+ else:
291
+ logger.warning(f"FAISS index {i} out of bounds for memory_items_list (len: {len(_memory_items_list)})")
292
+
293
+ logger.debug(f"Retrieved {len(results)} memories semantically for query: '{query[:50]}...'")
294
+ return results
295
  except Exception as e:
296
+ logger.error(f"Error retrieving memories semantically: {e}", exc_info=True)
297
+ return []
298
 
299
 
300
+ # --- Rule (Insight) Operations (Semantic) ---
301
+ def add_rule_entry(rule_text: str) -> tuple[bool, str]:
302
+ """Adds a rule if valid and not a duplicate. Updates backend and FAISS."""
303
+ global _rules_items_list, _faiss_rules_index
304
+ if not _initialized: initialize_memory_system()
305
+ if not _embedder or not _faiss_rules_index:
306
+ return False, "Rule system or embedder not initialized."
307
+
308
+ rule_text = rule_text.strip()
309
+ if not rule_text: return False, "Rule text cannot be empty."
310
+ if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\](.*)", rule_text, re.I|re.DOTALL):
311
+ return False, "Invalid rule format."
312
+ if rule_text in _rules_items_list:
313
+ return False, "duplicate"
314
 
 
 
 
 
 
315
  try:
316
+ embedding = _embedder.encode([rule_text], convert_to_tensor=False)
317
+ embedding_np = np.array(embedding, dtype=np.float32).reshape(1, -1)
318
+
319
+ if embedding_np.shape != (1, _dimension):
320
+ return False, "Rule embedding shape error."
321
+
322
+ _faiss_rules_index.add(embedding_np)
323
+ _rules_items_list.append(rule_text)
324
+ _rules_items_list.sort()
325
+
326
+ if STORAGE_BACKEND == "SQLITE" and sqlite3:
327
+ with _get_sqlite_connection() as conn:
328
+ conn.execute("INSERT OR IGNORE INTO rules (rule_text) VALUES (?)", (rule_text,))
329
+ conn.commit()
330
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
331
+ logger.info(f"Pushing {len(_rules_items_list)} rules to HF Hub: {HF_RULES_DATASET_REPO}")
332
+ Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
333
+
334
+ logger.info(f"Added rule. RAM: {len(_rules_items_list)}, FAISS: {_faiss_rules_index.ntotal}")
335
+ return True, "Rule added successfully."
336
  except Exception as e:
337
+ logger.error(f"Error adding rule entry: {e}", exc_info=True)
338
+ # Basic rollback if FAISS add succeeded
339
+ if rule_text in _rules_items_list and _faiss_rules_index.ntotal > 0: # Crude check
340
+ # A full rollback would involve rebuilding FAISS index from _rules_items_list before append.
341
+ # For simplicity, this is omitted here. State could be inconsistent on error.
342
+ pass
343
+ return False, f"Error adding rule: {e}"
344
 
345
+ def retrieve_rules_semantic(query: str, k: int = 5) -> list[str]:
346
+ """Retrieves k most relevant rules using semantic search."""
347
+ if not _initialized: initialize_memory_system()
348
+ if not _embedder or not _faiss_rules_index or _faiss_rules_index.ntotal == 0:
349
+ return []
350
+ try:
351
+ query_embedding = _embedder.encode([query], convert_to_tensor=False)
352
+ query_embedding_np = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
353
+
354
+ if query_embedding_np.shape[1] != _dimension: return []
355
 
356
+ distances, indices = _faiss_rules_index.search(query_embedding_np, min(k, _faiss_rules_index.ntotal))
357
+ results = [_rules_items_list[i] for i in indices[0] if 0 <= i < len(_rules_items_list)]
358
+ logger.debug(f"Retrieved {len(results)} rules semantically for query: '{query[:50]}...'")
359
+ return results
360
+ except Exception as e:
361
+ logger.error(f"Error retrieving rules semantically: {e}", exc_info=True)
362
+ return []
363
 
364
+ def remove_rule_entry(rule_text_to_delete: str) -> bool:
365
+ """Removes a rule from backend and rebuilds FAISS for rules."""
366
+ global _rules_items_list, _faiss_rules_index
367
+ if not _initialized: initialize_memory_system()
368
+ if not _embedder or not _faiss_rules_index: return False
369
+
370
+ rule_text_to_delete = rule_text_to_delete.strip()
371
+ if rule_text_to_delete not in _rules_items_list:
372
+ return False # Not found
373
 
374
  try:
375
+ _rules_items_list.remove(rule_text_to_delete)
376
+ _rules_items_list.sort() # Maintain sorted order
377
+
378
+ # Rebuild FAISS index for rules (simplest way to ensure consistency after removal)
379
+ new_faiss_rules_index = faiss.IndexFlatL2(_dimension)
380
+ if _rules_items_list:
381
+ embeddings = _embedder.encode(_rules_items_list, convert_to_tensor=False)
382
+ embeddings_np = np.array(embeddings, dtype=np.float32)
383
+ if embeddings_np.ndim == 2 and embeddings_np.shape[0] == len(_rules_items_list) and embeddings_np.shape[1] == _dimension:
384
+ new_faiss_rules_index.add(embeddings_np)
385
+ else: # Should not happen if list is consistent
386
+ logger.error("Error rebuilding FAISS for rules after removal: Embedding shape error. State might be inconsistent.")
387
+ # Attempt to revert _rules_items_list (add back the rule)
388
+ _rules_items_list.append(rule_text_to_delete)
389
+ _rules_items_list.sort()
390
+ return False # Indicate failure
391
+ _faiss_rules_index = new_faiss_rules_index
392
+
393
+ # Remove from persistent storage
394
+ if STORAGE_BACKEND == "SQLITE" and sqlite3:
395
+ with _get_sqlite_connection() as conn:
396
+ conn.execute("DELETE FROM rules WHERE rule_text = ?", (rule_text_to_delete,))
397
+ conn.commit()
398
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
399
+ Dataset.from_dict({"rule_text": list(_rules_items_list)}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
400
+
401
+ logger.info(f"Removed rule. RAM: {len(_rules_items_list)}, FAISS: {_faiss_rules_index.ntotal}")
402
  return True
403
  except Exception as e:
404
+ logger.error(f"Error removing rule entry: {e}", exc_info=True)
405
+ # Potential partial failure, state might be inconsistent.
406
  return False
407
 
408
+ # --- Utility functions to get all data (for UI display, etc.) ---
409
+ def get_all_rules_cached() -> list[str]:
410
+ if not _initialized: initialize_memory_system()
411
+ return list(_rules_items_list)
412
+
413
+ def get_all_memories_cached() -> list[dict]:
414
+ if not _initialized: initialize_memory_system()
415
+ # Convert JSON strings to dicts for easier use by UI
416
+ mem_dicts = []
417
+ for mem_json_str in _memory_items_list:
418
+ try: mem_dicts.append(json.loads(mem_json_str))
419
+ except: pass # Ignore parse errors for display
420
+ return mem_dicts
421
+
422
+ def clear_all_memory_data_backend() -> bool:
423
+ """Clears all memories from backend and resets in-memory FAISS/list."""
424
+ global _memory_items_list, _faiss_memory_index
425
+ if not _initialized: initialize_memory_system()
426
+
427
+ success = True
428
  try:
429
+ if STORAGE_BACKEND == "SQLITE" and sqlite3:
430
+ with _get_sqlite_connection() as conn: conn.execute("DELETE FROM memories"); conn.commit()
431
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
432
+ # Deleting from HF usually means pushing an empty dataset
433
+ Dataset.from_dict({"memory_json": []}).push_to_hub(HF_MEMORY_DATASET_REPO, token=HF_TOKEN, private=True)
434
+
435
+ _memory_items_list = []
436
+ if _faiss_memory_index: _faiss_memory_index.reset() # Clear FAISS index
437
+ logger.info("All memories cleared from backend and in-memory stores.")
438
  except Exception as e:
439
+ logger.error(f"Error clearing all memory data: {e}")
440
+ success = False
441
+ return success
442
 
443
+ def clear_all_rules_data_backend() -> bool:
444
+ """Clears all rules from backend and resets in-memory FAISS/list."""
445
+ global _rules_items_list, _faiss_rules_index
446
+ if not _initialized: initialize_memory_system()
447
+
448
+ success = True
449
  try:
450
+ if STORAGE_BACKEND == "SQLITE" and sqlite3:
451
+ with _get_sqlite_connection() as conn: conn.execute("DELETE FROM rules"); conn.commit()
452
+ elif STORAGE_BACKEND == "HF_DATASET" and HF_TOKEN and Dataset:
453
+ Dataset.from_dict({"rule_text": []}).push_to_hub(HF_RULES_DATASET_REPO, token=HF_TOKEN, private=True)
454
+
455
+ _rules_items_list = []
456
+ if _faiss_rules_index: _faiss_rules_index.reset()
457
+ logger.info("All rules cleared from backend and in-memory stores.")
458
  except Exception as e:
459
+ logger.error(f"Error clearing all rules data: {e}")
460
+ success = False
461
+ return success
462
+
463
+ # Optional: Function to save FAISS indices to disk (from ai-learn, if needed for persistence between app runs with RAM backend)
464
+ FAISS_MEMORY_PATH = os.path.join(os.getenv("FAISS_STORAGE_PATH", "app_data/faiss_indices"), "memory_index.faiss")
465
+ FAISS_RULES_PATH = os.path.join(os.getenv("FAISS_STORAGE_PATH", "app_data/faiss_indices"), "rules_index.faiss")
466
+
467
+ def save_faiss_indices_to_disk():
468
+ if not _initialized or not faiss: return
469
+
470
+ faiss_dir = os.path.dirname(FAISS_MEMORY_PATH)
471
+ if not os.path.exists(faiss_dir): os.makedirs(faiss_dir, exist_ok=True)
472
+
473
+ if _faiss_memory_index and _faiss_memory_index.ntotal > 0:
474
+ try:
475
+ faiss.write_index(_faiss_memory_index, FAISS_MEMORY_PATH)
476
+ logger.info(f"Memory FAISS index saved to disk ({_faiss_memory_index.ntotal} items).")
477
+ except Exception as e: logger.error(f"Error saving memory FAISS index: {e}")
478
+
479
+ if _faiss_rules_index and _faiss_rules_index.ntotal > 0:
480
+ try:
481
+ faiss.write_index(_faiss_rules_index, FAISS_RULES_PATH)
482
+ logger.info(f"Rules FAISS index saved to disk ({_faiss_rules_index.ntotal} items).")
483
+ except Exception as e: logger.error(f"Error saving rules FAISS index: {e}")
484
+
485
+ def load_faiss_indices_from_disk():
486
+ global _faiss_memory_index, _faiss_rules_index
487
+ if not _initialized or not faiss: return
488
+
489
+ if os.path.exists(FAISS_MEMORY_PATH) and _faiss_memory_index: # Check if index object exists
490
+ try:
491
+ logger.info(f"Loading memory FAISS index from {FAISS_MEMORY_PATH}...")
492
+ _faiss_memory_index = faiss.read_index(FAISS_MEMORY_PATH)
493
+ logger.info(f"Memory FAISS index loaded ({_faiss_memory_index.ntotal} items).")
494
+ # Consistency check: FAISS ntotal vs len(_memory_items_list)
495
+ if _faiss_memory_index.ntotal != len(_memory_items_list) and len(_memory_items_list) > 0:
496
+ logger.warning(f"Memory FAISS index count ({_faiss_memory_index.ntotal}) differs from loaded texts ({len(_memory_items_list)}). Consider rebuilding FAISS.")
497
+ except Exception as e: logger.error(f"Error loading memory FAISS index: {e}. Will use fresh index.")
498
+
499
+ if os.path.exists(FAISS_RULES_PATH) and _faiss_rules_index:
500
+ try:
501
+ logger.info(f"Loading rules FAISS index from {FAISS_RULES_PATH}...")
502
+ _faiss_rules_index = faiss.read_index(FAISS_RULES_PATH)
503
+ logger.info(f"Rules FAISS index loaded ({_faiss_rules_index.ntotal} items).")
504
+ if _faiss_rules_index.ntotal != len(_rules_items_list) and len(_rules_items_list) > 0:
505
+ logger.warning(f"Rules FAISS index count ({_faiss_rules_index.ntotal}) differs from loaded texts ({len(_rules_items_list)}). Consider rebuilding FAISS.")
506
+ except Exception as e: logger.error(f"Error loading rules FAISS index: {e}. Will use fresh index.")
507
+
508
+