| """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 |
| ); |
| """) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
