Upload jarvis/memory.py

jarvis/memory.py

@@ -0,0 +1,205 @@
+"""Memory layer: SQLite + ChromaDB for J.A.R.V.I.S."""
+import sqlite3
+import json
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Optional, Any
+
+import chromadb
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+
+from jarvis.config import MemoryConfig
+
+
+class SQLiteStore:
+    """Structured data: user state, open loops, preferences, mirror log."""
+
+    def __init__(self, db_path: str):
+        self.db_path = db_path
+        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
+        self._init_tables()
+
+    def _connect(self):
+        return sqlite3.connect(self.db_path)
+
+    def _init_tables(self):
+        with self._connect() as conn:
+            conn.executescript("""
+                CREATE TABLE IF NOT EXISTS user_state (
+                    key TEXT PRIMARY KEY,
+                    value TEXT,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                );
+                CREATE TABLE IF NOT EXISTS open_loops (
+                    id TEXT PRIMARY KEY,
+                    title TEXT,
+                    description TEXT,
+                    status TEXT DEFAULT 'open',
+                    priority INTEGER DEFAULT 5,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    last_nudged_at TIMESTAMP
+                );
+                CREATE TABLE IF NOT EXISTS mirror_check_log (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    turn_id TEXT,
+                    contracts TEXT,
+                    dials TEXT,
+                    drift_detected BOOLEAN,
+                    notes TEXT,
+                    checked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                );
+                CREATE TABLE IF NOT EXISTS exchanges (
+                    id TEXT PRIMARY KEY,
+                    user_message TEXT,
+                    jarvis_response TEXT,
+                    contract TEXT,
+                    dial INTEGER,
+                    archetype_mix TEXT,
+                    score REAL,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                );
+            """)
+
+    # ---- User State ----
+    def set_state(self, key: str, value: Any):
+        with self._connect() as conn:
+            conn.execute(
+                "INSERT INTO user_state (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value, updated_at=CURRENT_TIMESTAMP",
+                (key, json.dumps(value)),
+            )
+
+    def get_state(self, key: str, default: Any = None) -> Any:
+        with self._connect() as conn:
+            row = conn.execute("SELECT value FROM user_state WHERE key=?", (key,)).fetchone()
+        if row:
+            return json.loads(row[0])
+        return default
+
+    # ---- Open Loops ----
+    def add_loop(self, title: str, description: str = "", priority: int = 5) -> str:
+        lid = str(uuid.uuid4())
+        with self._connect() as conn:
+            conn.execute(
+                "INSERT INTO open_loops (id, title, description, priority) VALUES (?, ?, ?, ?)",
+                (lid, title, description, priority),
+            )
+        return lid
+
+    def list_loops(self, status: Optional[str] = None) -> List[Dict]:
+        with self._connect() as conn:
+            if status:
+                rows = conn.execute("SELECT * FROM open_loops WHERE status=? ORDER BY priority DESC", (status,)).fetchall()
+            else:
+                rows = conn.execute("SELECT * FROM open_loops ORDER BY priority DESC").fetchall()
+        cols = [d[0] for d in conn.execute("SELECT * FROM open_loops LIMIT 0").description]
+        return [dict(zip(cols, row)) for row in rows]
+
+    def update_loop(self, loop_id: str, **kwargs):
+        sets = ", ".join(f"{k}=?" for k in kwargs)
+        vals = list(kwargs.values()) + [loop_id]
+        with self._connect() as conn:
+            conn.execute(f"UPDATE open_loops SET {sets}, updated_at=CURRENT_TIMESTAMP WHERE id=?", vals)
+
+    # ---- Exchanges ----
+    def log_exchange(self, user_msg: str, response: str, contract: str, dial: int,
+                     archetype_mix: Dict[str, float], score: Optional[float] = None) -> str:
+        eid = str(uuid.uuid4())
+        with self._connect() as conn:
+            conn.execute(
+                "INSERT INTO exchanges (id, user_message, jarvis_response, contract, dial, archetype_mix, score) VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (eid, user_msg, response, contract, dial, json.dumps(archetype_mix), score),
+            )
+        return eid
+
+    def get_recent_exchanges(self, n: int = 5) -> List[Dict]:
+        with self._connect() as conn:
+            rows = conn.execute(
+                "SELECT * FROM exchanges ORDER BY timestamp DESC LIMIT ?", (n,)
+            ).fetchall()
+        if not rows:
+            return []
+        cols = [d[0] for d in conn.execute("SELECT * FROM exchanges LIMIT 0").description]
+        return [dict(zip(cols, row)) for row in reversed(rows)]
+
+
+class ChromaStore:
+    """Vector memory for semantic recall of conversation history."""
+
+    def __init__(self, persist_dir: str, embedding_model: str = "all-MiniLM-L6-v2"):
+        Path(persist_dir).mkdir(parents=True, exist_ok=True)
+        self.client = chromadb.PersistentClient(
+            path=persist_dir,
+            settings=Settings(anonymized_telemetry=False),
+        )
+        self.collection = self.client.get_or_create_collection("jarvis_memory")
+        self.embedder = SentenceTransformer(embedding_model)
+
+    def add(self, text: str, metadata: Optional[Dict] = None, doc_id: Optional[str] = None):
+        doc_id = doc_id or str(uuid.uuid4())
+        embedding = self.embedder.encode(text).tolist()
+        self.collection.add(
+            ids=[doc_id],
+            embeddings=[embedding],
+            documents=[text],
+            metadatas=[metadata or {}],
+        )
+
+    def query(self, query_text: str, n_results: int = 5) -> List[Dict]:
+        embedding = self.embedder.encode(query_text).tolist()
+        results = self.collection.query(
+            query_embeddings=[embedding],
+            n_results=n_results,
+            include=["documents", "metadatas", "distances"],
+        )
+        out = []
+        for i in range(len(results["ids"][0])):
+            out.append({
+                "id": results["ids"][0][i],
+                "document": results["documents"][0][i],
+                "metadata": results["metadatas"][0][i],
+                "distance": results["distances"][0][i],
+            })
+        return out
+
+
+class Memory:
+    """Unified memory interface."""
+
+    def __init__(self, config: MemoryConfig):
+        self.sqlite = SQLiteStore(config.sqlite_path)
+        self.chroma = None
+        if config.use_chroma:
+            self.chroma = ChromaStore(config.chroma_path, config.embedding_model)
+        self.max_context_turns = config.max_context_turns
+
+    def log_turn(self, user_msg: str, response: str, contract: str, dial: int,
+                 archetype_mix: Dict[str, float], score: Optional[float] = None):
+        eid = self.sqlite.log_exchange(user_msg, response, contract, dial, archetype_mix, score)
+        if self.chroma:
+            self.chroma.add(
+                text=f"User: {user_msg}\nJ.A.R.V.I.S.: {response}",
+                metadata={
+                    "exchange_id": eid,
+                    "contract": contract,
+                    "dial": dial,
+                    "timestamp": datetime.utcnow().isoformat(),
+                },
+                doc_id=eid,
+            )
+
+    def get_recent_context(self, n: Optional[int] = None) -> str:
+        n = n or self.max_context_turns
+        exchanges = self.sqlite.get_recent_exchanges(n)
+        lines = []
+        for ex in exchanges:
+            lines.append(f"User: {ex['user_message']}")
+            lines.append(f"J.A.R.V.I.S.: {ex['jarvis_response']}")
+        return "\n".join(lines)
+
+    def semantic_recall(self, query: str, n: int = 3) -> List[Dict]:
+        if self.chroma is None:
+            return []
+        return self.chroma.query(query, n_results=n)