Spaces:
Running on Zero
Running on Zero
| 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,)) | |