from __future__ import annotations import json import os import sqlite3 from pathlib import Path from typing import Any DB_PATH = Path(os.getenv("LEXI_DB_PATH", "lexi_v13.db")) SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS memories ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, source TEXT NOT NULL, target TEXT NOT NULL, kind TEXT NOT NULL, confidence REAL NOT NULL, evidence_count INTEGER NOT NULL DEFAULT 0, positive_evidence_count INTEGER NOT NULL DEFAULT 0, negative_evidence_count INTEGER NOT NULL DEFAULT 0, explicit_evidence_count INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL, aliases_json TEXT NOT NULL DEFAULT '[]', positive_contexts_json TEXT NOT NULL DEFAULT '[]', negative_contexts_json TEXT NOT NULL DEFAULT '[]', exceptions_json TEXT NOT NULL DEFAULT '[]', scope TEXT NOT NULL DEFAULT 'global', last_reason TEXT NOT NULL DEFAULT '', last_used_at TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE(user_id, source, target) ); CREATE INDEX IF NOT EXISTS idx_memories_user ON memories(user_id); CREATE INDEX IF NOT EXISTS idx_memories_status ON memories(user_id, status); CREATE TABLE IF NOT EXISTS memory_evidence ( id INTEGER PRIMARY KEY AUTOINCREMENT, memory_id TEXT NOT NULL, user_id TEXT NOT NULL, evidence_type TEXT NOT NULL, before_text TEXT NOT NULL, after_text TEXT NOT NULL, context TEXT NOT NULL DEFAULT '', note_title TEXT NOT NULL DEFAULT '', polarity TEXT NOT NULL DEFAULT 'positive', confidence REAL, created_at TEXT NOT NULL, FOREIGN KEY(memory_id) REFERENCES memories(id) ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS idx_memory_evidence_user ON memory_evidence(user_id); CREATE TABLE IF NOT EXISTS decision_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id TEXT NOT NULL, note_title TEXT NOT NULL, input_text TEXT NOT NULL, mode TEXT NOT NULL, decision_type TEXT NOT NULL, reason TEXT NOT NULL, trace_json TEXT NOT NULL, latency_ms REAL, input_tokens INTEGER, output_tokens INTEGER, estimated_cost_usd REAL, created_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_decision_log_user ON decision_log(user_id, created_at DESC); """ class MemoryStore: def __init__(self, path: Path | str = DB_PATH) -> None: self.path = Path(path) self.path.parent.mkdir(parents=True, exist_ok=True) self.initialize() def connect(self) -> sqlite3.Connection: connection = sqlite3.connect(self.path) connection.row_factory = sqlite3.Row connection.execute("PRAGMA foreign_keys = ON") return connection def initialize(self) -> None: with self.connect() as connection: connection.executescript(SCHEMA_SQL) columns = { row[1] for row in connection.execute("PRAGMA table_info(memories)").fetchall() } migrations = { "positive_contexts_json": "ALTER TABLE memories ADD COLUMN positive_contexts_json TEXT NOT NULL DEFAULT '[]'", "negative_contexts_json": "ALTER TABLE memories ADD COLUMN negative_contexts_json TEXT NOT NULL DEFAULT '[]'", "positive_evidence_count": "ALTER TABLE memories ADD COLUMN positive_evidence_count INTEGER NOT NULL DEFAULT 0", "negative_evidence_count": "ALTER TABLE memories ADD COLUMN negative_evidence_count INTEGER NOT NULL DEFAULT 0", "explicit_evidence_count": "ALTER TABLE memories ADD COLUMN explicit_evidence_count INTEGER NOT NULL DEFAULT 0", "exceptions_json": "ALTER TABLE memories ADD COLUMN exceptions_json TEXT NOT NULL DEFAULT '[]'", "last_used_at": "ALTER TABLE memories ADD COLUMN last_used_at TEXT", } for name, statement in migrations.items(): if name not in columns: connection.execute(statement) evidence_columns = { row[1] for row in connection.execute("PRAGMA table_info(memory_evidence)").fetchall() } evidence_migrations = { "note_title": "ALTER TABLE memory_evidence ADD COLUMN note_title TEXT NOT NULL DEFAULT ''", "polarity": "ALTER TABLE memory_evidence ADD COLUMN polarity TEXT NOT NULL DEFAULT 'positive'", } for name, statement in evidence_migrations.items(): if name not in evidence_columns: connection.execute(statement) def load_memories(self, user_id: str) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( "SELECT * FROM memories WHERE user_id = ? ORDER BY updated_at DESC", (user_id,), ).fetchall() result = [] for row in rows: item = dict(row) item["aliases"] = json.loads(item.pop("aliases_json", "[]")) item["positive_contexts"] = json.loads(item.pop("positive_contexts_json", "[]")) item["negative_contexts"] = json.loads(item.pop("negative_contexts_json", "[]")) item["exceptions"] = json.loads(item.pop("exceptions_json", "[]")) evidence_rows = connection.execute( """ SELECT evidence_type, before_text, after_text, context, note_title, polarity, confidence, created_at FROM memory_evidence WHERE memory_id = ? AND user_id = ? ORDER BY id DESC LIMIT 16 """, (item["id"], user_id), ).fetchall() item["evidence"] = [ { "type": evidence["evidence_type"], "before": evidence["before_text"], "after": evidence["after_text"], "context": evidence["context"], "note_title": evidence["note_title"], "polarity": evidence["polarity"], "confidence": evidence["confidence"], "timestamp": evidence["created_at"], } for evidence in reversed(evidence_rows) ] result.append(item) return result def upsert_memories(self, user_id: str, memories: list[dict[str, Any]]) -> None: with self.connect() as connection: for memory in memories: connection.execute( """ INSERT INTO memories ( id, user_id, source, target, kind, confidence, evidence_count, positive_evidence_count, negative_evidence_count, explicit_evidence_count, status, aliases_json, positive_contexts_json, negative_contexts_json, exceptions_json, scope, last_reason, last_used_at, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET source=excluded.source, target=excluded.target, kind=excluded.kind, confidence=excluded.confidence, evidence_count=excluded.evidence_count, positive_evidence_count=excluded.positive_evidence_count, negative_evidence_count=excluded.negative_evidence_count, explicit_evidence_count=excluded.explicit_evidence_count, status=excluded.status, aliases_json=excluded.aliases_json, positive_contexts_json=excluded.positive_contexts_json, negative_contexts_json=excluded.negative_contexts_json, exceptions_json=excluded.exceptions_json, scope=excluded.scope, last_reason=excluded.last_reason, last_used_at=excluded.last_used_at, updated_at=excluded.updated_at """, ( memory["id"], user_id, memory["source"], memory["target"], memory["kind"], memory["confidence"], memory.get("evidence_count", 0), memory.get("positive_evidence_count", 0), memory.get("negative_evidence_count", 0), memory.get("explicit_evidence_count", 0), memory["status"], json.dumps(memory.get("aliases", []), ensure_ascii=False), json.dumps(memory.get("positive_contexts", []), ensure_ascii=False), json.dumps(memory.get("negative_contexts", []), ensure_ascii=False), json.dumps(memory.get("exceptions", []), ensure_ascii=False), memory.get("scope", "global"), memory.get("last_reason", ""), memory.get("last_used_at"), memory["created_at"], memory["updated_at"], ), ) def save_evidence_for_memory( self, user_id: str, memory: dict[str, Any], before: str, after: str, evidence_type: str, polarity: str = "positive", note_title: str = "", context: str = "", confidence: float | None = None, timestamp: str | None = None, ) -> None: timestamp = timestamp or memory.get("updated_at") with self.connect() as connection: exists = connection.execute( """ SELECT 1 FROM memory_evidence WHERE memory_id = ? AND user_id = ? AND created_at = ? LIMIT 1 """, (memory["id"], user_id, timestamp), ).fetchone() if exists: return connection.execute( """ INSERT INTO memory_evidence ( memory_id, user_id, evidence_type, before_text, after_text, context, note_title, polarity, confidence, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( memory["id"], user_id, evidence_type, before, after, context, note_title, polarity, confidence, timestamp, ), ) def log_decision( self, user_id: str, note_title: str, input_text: str, mode: str, decision_type: str, reason: str, trace: dict[str, Any], metrics: dict[str, Any] | None = None, ) -> None: metrics = metrics or {} with self.connect() as connection: connection.execute( """ INSERT INTO decision_log ( user_id, note_title, input_text, mode, decision_type, reason, trace_json, latency_ms, input_tokens, output_tokens, estimated_cost_usd, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now')) """, ( user_id, note_title, input_text, mode, decision_type, reason, json.dumps(trace, ensure_ascii=False), metrics.get("latency_ms"), metrics.get("input_tokens"), metrics.get("output_tokens"), metrics.get("estimated_cost_usd"), ), ) def latest_decisions(self, user_id: str, limit: int = 8) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT * FROM decision_log WHERE user_id = ? ORDER BY id DESC LIMIT ? """, (user_id, limit), ).fetchall() result = [] for row in rows: item = dict(row) item["trace"] = json.loads(item.pop("trace_json")) result.append(item) return result def delete_memory(self, user_id: str, memory_id: str) -> None: with self.connect() as connection: connection.execute( "DELETE FROM memories WHERE id = ? AND user_id = ?", (memory_id, user_id), ) def reset_user(self, user_id: str) -> None: with self.connect() as connection: connection.execute("DELETE FROM decision_log WHERE user_id = ?", (user_id,)) connection.execute("DELETE FROM memory_evidence WHERE user_id = ?", (user_id,)) connection.execute("DELETE FROM memories WHERE user_id = ?", (user_id,))