File-Based and Database-Backed Long-Term Memory Stores

import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Protocol
 
 
def utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()
 
 
class MemoryStore(Protocol):
    def upsert(self, user_id: str, key: str, value: Any, *, namespace: str = "profile") -> None: ...
    def get(self, user_id: str, key: str, *, namespace: str = "profile") -> Any | None: ...
    def list(self, user_id: str, *, namespace: str | None = None) -> list[dict]: ...
    def delete(self, user_id: str, key: str, *, namespace: str = "profile") -> bool: ...
 
 
class JsonFileMemoryStore:
    def __init__(self, path: str | Path = "ltm.json") -> None:
        self.path = Path(path)
        if not self.path.exists():
            self.path.write_text('{"rows":[]}', encoding="utf-8")
 
    def _load(self) -> dict:
        return json.loads(self.path.read_text(encoding="utf-8"))
 
    def _save(self, data: dict) -> None:
        self.path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
 
    def upsert(self, user_id: str, key: str, value: Any, *, namespace: str = "profile") -> None:
        data = self._load()
        now = utcnow()
        for row in data["rows"]:
            if row["user_id"] == user_id and row["namespace"] == namespace and row["key"] == key:
                row["value"] = value
                row["updated_at"] = now
                self._save(data)
                return
        data["rows"].append({
            "user_id": user_id,
            "namespace": namespace,
            "key": key,
            "value": value,
            "created_at": now,
            "updated_at": now,
        })
        self._save(data)
 
    def get(self, user_id: str, key: str, *, namespace: str = "profile") -> Any | None:
        for row in self.list(user_id, namespace=namespace):
            if row["key"] == key:
                return row["value"]
        return None
 
    def list(self, user_id: str, *, namespace: str | None = None) -> list[dict]:
        rows = [r for r in self._load()["rows"] if r["user_id"] == user_id]
        if namespace is not None:
            rows = [r for r in rows if r["namespace"] == namespace]
        return rows
 
    def delete(self, user_id: str, key: str, *, namespace: str = "profile") -> bool:
        data = self._load()
        before = len(data["rows"])
        data["rows"] = [
            r for r in data["rows"]
            if not (r["user_id"] == user_id and r["namespace"] == namespace and r["key"] == key)
        ]
        self._save(data)
        return len(data["rows"]) < before
 
 
class SqliteMemoryStore:
    def __init__(self, path: str = "ltm.db") -> None:
        self.conn = sqlite3.connect(path)
        self.conn.row_factory = sqlite3.Row
        self.conn.execute(
            """
            CREATE TABLE IF NOT EXISTS memories (
              user_id TEXT NOT NULL,
              namespace TEXT NOT NULL,
              key TEXT NOT NULL,
              value_json TEXT NOT NULL,
              created_at TEXT NOT NULL,
              updated_at TEXT NOT NULL,
              PRIMARY KEY (user_id, namespace, key)
            )
            """
        )
        self.conn.commit()
 
    def upsert(self, user_id: str, key: str, value: Any, *, namespace: str = "profile") -> None:
        now = utcnow()
        self.conn.execute(
            """
            INSERT INTO memories(user_id, namespace, key, value_json, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(user_id, namespace, key) DO UPDATE SET
              value_json=excluded.value_json,
              updated_at=excluded.updated_at
            """,
            (user_id, namespace, key, json.dumps(value), now, now),
        )
        self.conn.commit()
 
    def get(self, user_id: str, key: str, *, namespace: str = "profile") -> Any | None:
        cur = self.conn.execute(
            "SELECT value_json FROM memories WHERE user_id=? AND namespace=? AND key=?",
            (user_id, namespace, key),
        )
        row = cur.fetchone()
        return json.loads(row["value_json"]) if row else None
 
    def list(self, user_id: str, *, namespace: str | None = None) -> list[dict]:
        if namespace is None:
            cur = self.conn.execute(
                "SELECT user_id, namespace, key, value_json, created_at, updated_at FROM memories WHERE user_id=?",
                (user_id,),
            )
        else:
            cur = self.conn.execute(
                "SELECT user_id, namespace, key, value_json, created_at, updated_at FROM memories WHERE user_id=? AND namespace=?",
                (user_id, namespace),
            )
        return [
            {
                "user_id": r["user_id"],
                "namespace": r["namespace"],
                "key": r["key"],
                "value": json.loads(r["value_json"]),
                "created_at": r["created_at"],
                "updated_at": r["updated_at"],
            }
            for r in cur.fetchall()
        ]
 
    def delete(self, user_id: str, key: str, *, namespace: str = "profile") -> bool:
        cur = self.conn.execute(
            "DELETE FROM memories WHERE user_id=? AND namespace=? AND key=?",
            (user_id, namespace, key),
        )
        self.conn.commit()
        return cur.rowcount > 0
 
 
# Same application code against either backend:
store: MemoryStore = SqliteMemoryStore()
store.upsert("user-42", "preferred_language", "Portuguese")
store.upsert("user-42", "stack", {"language": "python", "framework": "fastapi"}, namespace="project")
print(store.get("user-42", "preferred_language"))
print(store.list("user-42"))