Spaces:
Running on Zero
Running on Zero
File size: 13,544 Bytes
9fca8b4 4afb0f2 9fca8b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 | 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,))
|