File size: 7,120 Bytes
0be710f | 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 | import json
import sqlite3
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional
DB_PATH = Path("/data/memory.db")
LOCK = threading.Lock()
def _ensure_db():
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
conn.executescript(
"""
CREATE TABLE IF NOT EXISTS user_facts (
id INTEGER PRIMARY KEY,
category TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
confidence REAL DEFAULT 1.0,
source TEXT DEFAULT 'user_stated',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(category, key)
);
CREATE TABLE IF NOT EXISTS chat_history (
id INTEGER PRIMARY KEY,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT,
model TEXT,
tool_call TEXT,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_chat_history_session ON chat_history(session_id, timestamp);
CREATE INDEX IF NOT EXISTS idx_user_facts_category ON user_facts(category);
"""
)
conn.commit()
conn.close()
def _connect():
_ensure_db()
return sqlite3.connect(str(DB_PATH), check_same_thread=False)
def store_fact(category: str, key: str, value: str, confidence: float = 1.0, source: str = "user_stated") -> Dict[str, Any]:
with LOCK:
conn = _connect()
try:
conn.execute(
"""
INSERT INTO user_facts (category, key, value, confidence, source, updated_at)
VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(category, key) DO UPDATE SET
value=excluded.value,
confidence=excluded.confidence,
source=excluded.source,
updated_at=CURRENT_TIMESTAMP
""",
(category, key, value, confidence, source),
)
conn.commit()
return {"ok": True, "category": category, "key": key, "value": value}
finally:
conn.close()
def recall_facts(query: Optional[str] = None, category: Optional[str] = None, limit: int = 20) -> List[Dict[str, Any]]:
with LOCK:
conn = _connect()
try:
cur = conn.cursor()
if category:
cur.execute(
"""
SELECT category, key, value, confidence, source, updated_at
FROM user_facts
WHERE category = ?
ORDER BY updated_at DESC
LIMIT ?
""",
(category, limit),
)
elif query:
like = f"%{query}%"
cur.execute(
"""
SELECT category, key, value, confidence, source, updated_at
FROM user_facts
WHERE key LIKE ? OR value LIKE ?
ORDER BY updated_at DESC
LIMIT ?
""",
(like, like, limit),
)
else:
cur.execute(
"""
SELECT category, key, value, confidence, source, updated_at
FROM user_facts
ORDER BY updated_at DESC
LIMIT ?
""",
(limit,),
)
rows = cur.fetchall()
return [
{
"category": r[0],
"key": r[1],
"value": r[2],
"confidence": r[3],
"source": r[4],
"updated_at": r[5],
}
for r in rows
]
finally:
conn.close()
def update_fact(category: str, key: str, value: str, confidence: Optional[float] = None) -> Dict[str, Any]:
with LOCK:
conn = _connect()
try:
params = [value]
sql = "UPDATE user_facts SET value = ?, updated_at = CURRENT_TIMESTAMP"
if confidence is not None:
sql += ", confidence = ?"
params.append(confidence)
sql += " WHERE category = ? AND key = ?"
params.extend([category, key])
cur = conn.cursor()
cur.execute(sql, params)
conn.commit()
if cur.rowcount == 0:
return {"ok": False, "error": "Fact not found", "category": category, "key": key}
return {"ok": True, "category": category, "key": key, "value": value}
finally:
conn.close()
def forget_fact(category: str, key: str) -> Dict[str, Any]:
with LOCK:
conn = _connect()
try:
cur = conn.cursor()
cur.execute("DELETE FROM user_facts WHERE category = ? AND key = ?", (category, key))
conn.commit()
return {"ok": True, "deleted": cur.rowcount, "category": category, "key": key}
finally:
conn.close()
def append_chat(session_id: str, role: str, content: Optional[str] = None, model: Optional[str] = None, tool_call: Optional[str] = None) -> Dict[str, Any]:
with LOCK:
conn = _connect()
try:
conn.execute(
"INSERT INTO chat_history (session_id, role, content, model, tool_call) VALUES (?, ?, ?, ?, ?)",
(session_id, role, content, model, tool_call),
)
conn.commit()
return {"ok": True, "session_id": session_id, "role": role}
finally:
conn.close()
def load_chat(session_id: str, limit: int = 100) -> List[Dict[str, Any]]:
with LOCK:
conn = _connect()
try:
cur = conn.cursor()
cur.execute(
"""
SELECT role, content, model, tool_call, timestamp
FROM chat_history
WHERE session_id = ?
ORDER BY timestamp ASC
LIMIT ?
""",
(session_id, limit),
)
return [
{"role": r[0], "content": r[1], "model": r[2], "tool_call": r[3], "timestamp": r[4]}
for r in cur.fetchall()
]
finally:
conn.close()
def build_memory_context() -> str:
"""Return a concise memory block for injection into the system prompt."""
facts = recall_facts(limit=30)
if not facts:
return ""
lines = ["User facts you should remember:"]
for f in facts:
lines.append(f"- [{f['category']}] {f['key']}: {f['value']} (confidence {f['confidence']:.1f})")
return "\n".join(lines)
def export_memory() -> Dict[str, Any]:
return {"facts": recall_facts(limit=10000), "exported_at": time.strftime("%Y-%m-%d %H:%M:%S")}
|