import os import sqlite3 import json import threading import time from contextlib import contextmanager import numpy as np from typing import List, Dict, Any, Optional from src.memory.working import WorkingMemory DEFAULT_DB_PATH = os.path.join("diagnostics", "memory_store.db") class PersistentMemoryManager: """ Multi-store persistent memory framework surviving process restarts. Manages Working Memory, Episodic Memory, Semantic Associative Memory, Skill Memory, and Dream Replay Memory. """ def __init__(self, db_path: str = DEFAULT_DB_PATH): self.db_path = db_path self._lock = threading.RLock() os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True) self.working = WorkingMemory(capacity=7) self._init_db() @contextmanager def _locked(self): self._lock.acquire() try: yield finally: self._lock.release() def _init_db(self): with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() # 1. Episodic memory table cursor.execute(""" CREATE TABLE IF NOT EXISTS episodic_memory ( id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp REAL, step INTEGER, observation TEXT, action TEXT, reward REAL, prediction_error REAL, outcome TEXT ) """) # 2. Semantic associative memory table cursor.execute(""" CREATE TABLE IF NOT EXISTS semantic_memory ( id INTEGER PRIMARY KEY AUTOINCREMENT, concept TEXT UNIQUE, description TEXT, embedding BLOB, associations TEXT ) """) # 3. Skill memory table cursor.execute(""" CREATE TABLE IF NOT EXISTS skill_memory ( skill_name TEXT PRIMARY KEY, tool_sequence TEXT, success_count INTEGER, attempt_count INTEGER, last_used REAL ) """) # 4. Dream & replay memory table cursor.execute(""" CREATE TABLE IF NOT EXISTS dream_memory ( id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp REAL, seed INTEGER, base_episode_id INTEGER, simulated_action TEXT, counterfactual_reward REAL, insight TEXT ) """) # 5. Experiment history table cursor.execute(""" CREATE TABLE IF NOT EXISTS experiment_history ( experiment_id TEXT PRIMARY KEY, timestamp REAL, seed INTEGER, config TEXT, results TEXT ) """) conn.commit() # --- Episodic Operations --- def record_episode( self, step: int, observation: Any, action: str, reward: float, prediction_error: float, outcome: Any ) -> int: with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO episodic_memory (timestamp, step, observation, action, reward, prediction_error, outcome) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( time.time(), step, json.dumps(observation) if not isinstance(observation, str) else observation, action, reward, prediction_error, json.dumps(outcome) if not isinstance(outcome, str) else outcome )) conn.commit() return cursor.lastrowid def get_recent_episodes(self, limit: int = 10) -> List[Dict[str, Any]]: with self._locked(), sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM episodic_memory ORDER BY id DESC LIMIT ?", (limit,)) rows = cursor.fetchall() return [dict(r) for r in rows] # --- Semantic Associative Operations --- def store_concept(self, concept: str, description: str, embedding: np.ndarray, associations: Optional[Dict] = None): emb_blob = np.asarray(embedding, dtype=np.float32).tobytes() assoc_str = json.dumps(associations or {}) with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO semantic_memory (concept, description, embedding, associations) VALUES (?, ?, ?, ?) ON CONFLICT(concept) DO UPDATE SET description=excluded.description, embedding=excluded.embedding, associations=excluded.associations """, (concept, description, emb_blob, assoc_str)) conn.commit() def query_semantic(self, query_emb: np.ndarray, top_k: int = 3) -> List[Dict[str, Any]]: query_vec = np.asarray(query_emb, dtype=np.float32).flatten() norm_q = np.linalg.norm(query_vec) + 1e-7 with self._locked(), sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT id, concept, description, embedding, associations FROM semantic_memory") rows = cursor.fetchall() scored = [] for r in rows: emb = np.frombuffer(r["embedding"], dtype=np.float32) sim = float(np.dot(query_vec, emb) / (norm_q * (np.linalg.norm(emb) + 1e-7))) scored.append({ "concept": r["concept"], "description": r["description"], "similarity": sim, "associations": json.loads(r["associations"]) }) scored.sort(key=lambda x: x["similarity"], reverse=True) return scored[:top_k] # --- Skill Memory Operations --- def record_skill(self, skill_name: str, tool_sequence: List[str], success: bool): with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute("SELECT success_count, attempt_count FROM skill_memory WHERE skill_name=?", (skill_name,)) row = cursor.fetchone() if row: s_count = row[0] + (1 if success else 0) a_count = row[1] + 1 cursor.execute(""" UPDATE skill_memory SET success_count=?, attempt_count=?, last_used=?, tool_sequence=? WHERE skill_name=? """, (s_count, a_count, time.time(), json.dumps(tool_sequence), skill_name)) else: cursor.execute(""" INSERT INTO skill_memory (skill_name, tool_sequence, success_count, attempt_count, last_used) VALUES (?, ?, ?, ?, ?) """, (skill_name, json.dumps(tool_sequence), 1 if success else 0, 1, time.time())) conn.commit() def get_skills(self) -> List[Dict[str, Any]]: with self._locked(), sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM skill_memory ORDER BY success_count DESC") return [ { "skill_name": r["skill_name"], "tool_sequence": json.loads(r["tool_sequence"]), "success_count": r["success_count"], "attempt_count": r["attempt_count"], "success_rate": round(r["success_count"] / max(1, r["attempt_count"]), 3) } for r in cursor.fetchall() ] # --- Dream & Replay Operations --- def record_dream( self, seed: int = 42, base_episode_id: int = 1, simulated_action: str = "explore", counterfactual_reward: float = 0.5, insight: str = "", **kwargs ) -> int: ep_id = kwargs.get("episode_id", base_episode_id) reward_val = kwargs.get("hypothetical_reward", counterfactual_reward) insight_str = kwargs.get("consolidation_insight", insight) seed_val = kwargs.get("seed", seed) action_val = kwargs.get("simulated_action", simulated_action) with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute(""" INSERT INTO dream_memory (timestamp, seed, base_episode_id, simulated_action, counterfactual_reward, insight) VALUES (?, ?, ?, ?, ?, ?) """, (time.time(), seed_val, ep_id, action_val, reward_val, insight_str)) conn.commit() return cursor.lastrowid def get_recent_dreams(self, limit: int = 10) -> List[Dict[str, Any]]: with self._locked(), sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.cursor() cursor.execute("SELECT * FROM dream_memory ORDER BY id DESC LIMIT ?", (limit,)) return [dict(r) for r in cursor.fetchall()] # --- Experiment History Operations --- def record_experiment(self, experiment_id: str, seed: int, config: Dict, results: Dict): with self._locked(), sqlite3.connect(self.db_path) as conn: cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO experiment_history (experiment_id, timestamp, seed, config, results) VALUES (?, ?, ?, ?, ?) """, (experiment_id, time.time(), seed, json.dumps(config), json.dumps(results))) conn.commit()