from __future__ import annotations import re import sqlite3 import uuid from contextlib import contextmanager from datetime import UTC, datetime from pathlib import Path MEMORY_PATTERNS = ( (re.compile(r"\bmy name is ([a-zA-Z][a-zA-Z\s'-]{1,40})", re.I), "Name"), (re.compile(r"\bi am allergic to ([^.!,;\n]+)", re.I), "Allergy"), (re.compile(r"\bi have ([^.!,;\n]+)", re.I), "Condition"), (re.compile(r"\bi am taking ([^.!,;\n]+)", re.I), "Medication"), (re.compile(r"\bmy age is (\d{1,3})\b", re.I), "Age"), ) class MemoryStore: def __init__(self, database_path: Path) -> None: self.database_path = Path(database_path) self.database_path.parent.mkdir(parents=True, exist_ok=True) self._initialize() @contextmanager def _connect(self): connection = sqlite3.connect(self.database_path) connection.row_factory = sqlite3.Row try: yield connection connection.commit() finally: connection.close() def _initialize(self) -> None: with self._connect() as connection: connection.executescript( """ CREATE TABLE IF NOT EXISTS users ( user_id TEXT PRIMARY KEY, created_at TEXT NOT NULL, last_seen_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id TEXT NOT NULL, role TEXT NOT NULL CHECK(role IN ('user', 'assistant')), content TEXT NOT NULL, created_at TEXT NOT NULL, FOREIGN KEY(user_id) REFERENCES users(user_id) ); CREATE TABLE IF NOT EXISTS memory_notes ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id TEXT NOT NULL, note TEXT NOT NULL, created_at TEXT NOT NULL, UNIQUE(user_id, note), FOREIGN KEY(user_id) REFERENCES users(user_id) ); """ ) def create_user(self) -> dict[str, str]: user_id = f"MB-{uuid.uuid4().hex[:12].upper()}" now = datetime.now(UTC).isoformat() with self._connect() as connection: connection.execute( "INSERT INTO users (user_id, created_at, last_seen_at) VALUES (?, ?, ?)", (user_id, now, now), ) return {"user_id": user_id, "created_at": now} def ensure_user(self, user_id: str) -> dict[str, str] | None: with self._connect() as connection: row = connection.execute( "SELECT user_id, created_at FROM users WHERE user_id = ?", (user_id,), ).fetchone() if not row: return None connection.execute( "UPDATE users SET last_seen_at = ? WHERE user_id = ?", (datetime.now(UTC).isoformat(), user_id), ) return {"user_id": row["user_id"], "created_at": row["created_at"]} def add_message(self, user_id: str, role: str, content: str) -> None: timestamp = datetime.now(UTC).isoformat() with self._connect() as connection: connection.execute( "INSERT INTO messages (user_id, role, content, created_at) VALUES (?, ?, ?, ?)", (user_id, role, content, timestamp), ) connection.execute( "UPDATE users SET last_seen_at = ? WHERE user_id = ?", (timestamp, user_id), ) if role == "user": self._extract_memory_notes(user_id, content) def _extract_memory_notes(self, user_id: str, content: str) -> None: extracted: set[str] = set() normalized = " ".join(content.split()) for pattern, label in MEMORY_PATTERNS: for match in pattern.findall(normalized): value = match.strip(" .") if 2 <= len(value) <= 120: extracted.add(f"{label}: {value}") if not extracted: return now = datetime.now(UTC).isoformat() with self._connect() as connection: connection.executemany( "INSERT OR IGNORE INTO memory_notes (user_id, note, created_at) VALUES (?, ?, ?)", [(user_id, note, now) for note in sorted(extracted)], ) def get_history(self, user_id: str, limit: int | None = None) -> list[dict[str, str]]: query = """ SELECT role, content, created_at FROM messages WHERE user_id = ? ORDER BY id DESC """ params: list[str | int] = [user_id] if limit is not None: query += " LIMIT ?" params.append(limit) with self._connect() as connection: rows = connection.execute(query, params).fetchall() return [dict(row) for row in reversed(rows)] def get_memory_notes(self, user_id: str, limit: int = 10) -> list[str]: with self._connect() as connection: rows = connection.execute( """ SELECT note FROM memory_notes WHERE user_id = ? ORDER BY id DESC LIMIT ? """, (user_id, limit), ).fetchall() return [row["note"] for row in rows]