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