memorytrace 0.1.0__py3-none-any.whl

engram/storage/sqlite_store.py

@@ -0,0 +1,615 @@
+"""SQLite + FTS5 storage backend — the heart of Engram."""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import threading
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from engram.exceptions import EntityAlreadyExistsError, EntityNotFoundError, StorageError
+from engram.models.entity import Entity, EntityState, Tier
+from engram.models.fact import Fact, FactStatus
+from engram.models.relation import Relation
+from engram.models.session import Session, SessionEvent
+from engram.models.source import Source, SourceType
+
+SCHEMA_VERSION = "1"
+
+_SCHEMA_SQL = """
+CREATE TABLE IF NOT EXISTS _meta (
+    key TEXT PRIMARY KEY,
+    value TEXT
+);
+
+CREATE TABLE IF NOT EXISTS entities (
+    id TEXT PRIMARY KEY,
+    name TEXT NOT NULL UNIQUE,
+    entity_type TEXT NOT NULL DEFAULT 'person',
+    tier TEXT NOT NULL DEFAULT 'recall' CHECK(tier IN ('core','recall','archival')),
+    summary TEXT DEFAULT '',
+    aliases TEXT DEFAULT '[]',
+    state TEXT DEFAULT '{}',
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL,
+    access_count INTEGER DEFAULT 0,
+    last_accessed TEXT
+);
+
+CREATE TABLE IF NOT EXISTS facts (
+    id TEXT PRIMARY KEY,
+    entity_id TEXT NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+    subject TEXT NOT NULL,
+    predicate TEXT NOT NULL,
+    object TEXT NOT NULL,
+    raw_text TEXT NOT NULL,
+    source_type TEXT NOT NULL,
+    source_author TEXT DEFAULT '',
+    source_channel TEXT DEFAULT '',
+    source_timestamp TEXT,
+    source_url TEXT,
+    source_session_id TEXT,
+    source_confidence REAL DEFAULT 1.0,
+    confidence REAL NOT NULL DEFAULT 0.5,
+    status TEXT NOT NULL DEFAULT 'unverified',
+    valid_from TEXT,
+    valid_to TEXT,
+    superseded_by TEXT,
+    created_at TEXT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS relations (
+    id TEXT PRIMARY KEY,
+    from_entity_id TEXT NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+    to_entity_id TEXT NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+    relation_type TEXT NOT NULL,
+    metadata TEXT DEFAULT '{}',
+    valid_from TEXT,
+    valid_to TEXT,
+    created_at TEXT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS sessions (
+    session_id TEXT PRIMARY KEY,
+    agent_id TEXT NOT NULL,
+    started_at TEXT NOT NULL,
+    ended_at TEXT,
+    parent_session_id TEXT,
+    entities_accessed TEXT DEFAULT '[]',
+    entities_modified TEXT DEFAULT '[]',
+    facts_added TEXT DEFAULT '[]',
+    summary TEXT,
+    metadata TEXT DEFAULT '{}'
+);
+
+CREATE TABLE IF NOT EXISTS session_events (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    session_id TEXT NOT NULL REFERENCES sessions(session_id),
+    event_type TEXT NOT NULL,
+    target TEXT,
+    detail TEXT DEFAULT '{}',
+    timestamp TEXT NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS conflicts (
+    id TEXT PRIMARY KEY,
+    existing_fact_id TEXT,
+    new_fact_id TEXT,
+    conflict_type TEXT NOT NULL,
+    suggested_resolution TEXT,
+    resolved_at TEXT,
+    resolved_by TEXT,
+    resolution TEXT
+);
+
+CREATE INDEX IF NOT EXISTS idx_entities_name ON entities(name);
+CREATE INDEX IF NOT EXISTS idx_entities_tier ON entities(tier);
+CREATE INDEX IF NOT EXISTS idx_facts_entity ON facts(entity_id);
+CREATE INDEX IF NOT EXISTS idx_facts_status ON facts(status);
+CREATE INDEX IF NOT EXISTS idx_facts_predicate ON facts(entity_id, predicate);
+CREATE INDEX IF NOT EXISTS idx_relations_from ON relations(from_entity_id);
+CREATE INDEX IF NOT EXISTS idx_relations_to ON relations(to_entity_id);
+CREATE INDEX IF NOT EXISTS idx_sessions_agent ON sessions(agent_id, started_at);
+"""
+
+_FTS_SQL = """
+CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
+    entity_id,
+    entity_name,
+    entity_type,
+    fact_text,
+    state_text,
+    summary,
+    tokenize='unicode61 remove_diacritics 2'
+);
+"""
+
+
+def _dt_to_str(dt: Optional[datetime]) -> Optional[str]:
+    return dt.isoformat() if dt else None
+
+
+def _str_to_dt(s: Optional[str]) -> Optional[datetime]:
+    return datetime.fromisoformat(s) if s else None
+
+
+class SQLiteStorage:
+    """SQLite + FTS5 storage backend with WAL mode and atomic writes."""
+
+    def __init__(self, db_path: Path):
+        self.db_path = Path(db_path).resolve()
+        self._local = threading.local()
+
+    @property
+    def _conn(self) -> sqlite3.Connection:
+        if not hasattr(self._local, "conn") or self._local.conn is None:
+            self._local.conn = self._create_connection()
+        return self._local.conn
+
+    def _create_connection(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA foreign_keys=ON")
+        conn.execute("PRAGMA busy_timeout=5000")
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def initialize(self) -> None:
+        """Create tables and FTS index if they don't exist."""
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        is_new = not self.db_path.exists()
+        conn = self._conn
+        conn.executescript(_SCHEMA_SQL)
+        conn.executescript(_FTS_SQL)
+        # Set schema version if not exists
+        conn.execute(
+            "INSERT OR IGNORE INTO _meta (key, value) VALUES (?, ?)",
+            ("schema_version", SCHEMA_VERSION),
+        )
+        conn.commit()
+        # Restrict DB file permissions to owner only (0o600)
+        if is_new and self.db_path.exists():
+            try:
+                os.chmod(self.db_path, 0o600)
+            except OSError:
+                pass  # Windows or restricted filesystem
+
+    def close(self) -> None:
+        if hasattr(self._local, "conn") and self._local.conn is not None:
+            self._local.conn.close()
+            self._local.conn = None
+
+    # ── Entity CRUD ──
+
+    def create_entity(self, entity: Entity) -> Entity:
+        now = datetime.now()
+        entity.created_at = now
+        entity.updated_at = now
+        try:
+            self._conn.execute(
+                """INSERT INTO entities
+                   (id, name, entity_type, tier, summary, aliases, state, created_at, updated_at, access_count, last_accessed)
+                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                (
+                    entity.id, entity.name, entity.entity_type, entity.tier.value,
+                    entity.summary, json.dumps(entity.aliases, ensure_ascii=False),
+                    json.dumps(entity.state.to_dict(), ensure_ascii=False),
+                    _dt_to_str(entity.created_at), _dt_to_str(entity.updated_at),
+                    entity.access_count, _dt_to_str(entity.last_accessed),
+                ),
+            )
+            # Add to FTS index
+            self._index_entity(entity)
+            self._conn.commit()
+        except sqlite3.IntegrityError:
+            self._conn.rollback()
+            raise EntityAlreadyExistsError(f"Entity '{entity.name}' already exists")
+        except Exception:
+            self._conn.rollback()
+            raise
+        return entity
+
+    def get_entity(self, entity_id: str) -> Optional[Entity]:
+        row = self._conn.execute(
+            "SELECT * FROM entities WHERE id = ?", (entity_id,)
+        ).fetchone()
+        return self._row_to_entity(row) if row else None
+
+    def get_entity_by_name(self, name: str) -> Optional[Entity]:
+        # Primary name lookup
+        row = self._conn.execute(
+            "SELECT * FROM entities WHERE name = ? COLLATE NOCASE", (name,)
+        ).fetchone()
+        if row:
+            return self._row_to_entity(row)
+        # Alias lookup (search within JSON array)
+        row = self._conn.execute(
+            "SELECT * FROM entities WHERE aliases LIKE ? COLLATE NOCASE",
+            (f'%"{name}"%',),
+        ).fetchone()
+        return self._row_to_entity(row) if row else None
+
+    def update_entity(self, entity: Entity) -> Entity:
+        entity.updated_at = datetime.now()
+        try:
+            self._conn.execute(
+                """UPDATE entities SET
+                   name=?, entity_type=?, tier=?, summary=?, aliases=?, state=?,
+                   updated_at=?, access_count=?, last_accessed=?
+                   WHERE id=?""",
+                (
+                    entity.name, entity.entity_type, entity.tier.value,
+                    entity.summary, json.dumps(entity.aliases, ensure_ascii=False),
+                    json.dumps(entity.state.to_dict(), ensure_ascii=False),
+                    _dt_to_str(entity.updated_at), entity.access_count,
+                    _dt_to_str(entity.last_accessed), entity.id,
+                ),
+            )
+            self._reindex_entity(entity)
+            self._conn.commit()
+        except Exception:
+            self._conn.rollback()
+            raise
+        return entity
+
+    def delete_entity(self, entity_id: str) -> bool:
+        try:
+            cursor = self._conn.execute("DELETE FROM entities WHERE id = ?", (entity_id,))
+            self._conn.execute(
+                "DELETE FROM memory_fts WHERE entity_id = ?", (entity_id,)
+            )
+            self._conn.commit()
+            return cursor.rowcount > 0
+        except Exception:
+            self._conn.rollback()
+            raise
+
+    def list_entities(
+        self,
+        tier: Optional[Tier] = None,
+        entity_type: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[Entity]:
+        query = "SELECT * FROM entities WHERE 1=1"
+        params: list = []
+        if tier:
+            query += " AND tier = ?"
+            params.append(tier.value)
+        if entity_type:
+            query += " AND entity_type = ?"
+            params.append(entity_type)
+        query += " ORDER BY updated_at DESC LIMIT ? OFFSET ?"
+        params.extend([limit, offset])
+        rows = self._conn.execute(query, params).fetchall()
+        return [self._row_to_entity(r) for r in rows]
+
+    # ── Fact CRUD ──
+
+    def add_fact(self, fact: Fact, reindex: bool = True) -> Fact:
+        fact.created_at = datetime.now()
+        try:
+            self._conn.execute(
+                """INSERT INTO facts
+                   (id, entity_id, subject, predicate, object, raw_text,
+                    source_type, source_author, source_channel, source_timestamp,
+                    source_url, source_session_id, source_confidence,
+                    confidence, status,
+                    valid_from, valid_to, superseded_by, created_at)
+                   VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
+                (
+                    fact.id, fact.entity_id, fact.subject, fact.predicate, fact.object,
+                    fact.raw_text, fact.source.type.value, fact.source.author,
+                    fact.source.channel, _dt_to_str(fact.source.timestamp),
+                    fact.source.url, fact.source.session_id,
+                    fact.source.confidence, fact.confidence, fact.status.value,
+                    _dt_to_str(fact.valid_from), _dt_to_str(fact.valid_to),
+                    fact.superseded_by, _dt_to_str(fact.created_at),
+                ),
+            )
+        except sqlite3.IntegrityError as e:
+            self._conn.rollback()
+            raise StorageError(f"Failed to add fact: {e}") from e
+        except Exception:
+            self._conn.rollback()
+            raise
+        # Update FTS index with new fact text (skippable for batch operations)
+        if reindex:
+            try:
+                entity = self.get_entity(fact.entity_id)
+                if entity:
+                    self._reindex_entity(entity)
+            except Exception:
+                self._conn.rollback()
+                raise
+        self._conn.commit()
+        return fact
+
+    def get_facts(
+        self, entity_id: str, status: Optional[FactStatus] = None
+    ) -> list[Fact]:
+        if status:
+            rows = self._conn.execute(
+                "SELECT * FROM facts WHERE entity_id = ? AND status = ? ORDER BY created_at DESC",
+                (entity_id, status.value),
+            ).fetchall()
+        else:
+            rows = self._conn.execute(
+                "SELECT * FROM facts WHERE entity_id = ? ORDER BY created_at DESC",
+                (entity_id,),
+            ).fetchall()
+        return [self._row_to_fact(r) for r in rows]
+
+    def get_current_facts(self, entity_id: str) -> list[Fact]:
+        rows = self._conn.execute(
+            """SELECT * FROM facts
+               WHERE entity_id = ? AND valid_to IS NULL AND superseded_by IS NULL
+               AND status NOT IN ('expired', 'retracted')
+               ORDER BY created_at DESC""",
+            (entity_id,),
+        ).fetchall()
+        return [self._row_to_fact(r) for r in rows]
+
+    def update_fact(self, fact: Fact) -> Fact:
+        self._conn.execute(
+            """UPDATE facts SET
+               subject=?, predicate=?, object=?, raw_text=?,
+               confidence=?, status=?, valid_from=?, valid_to=?,
+               superseded_by=?
+               WHERE id=?""",
+            (
+                fact.subject, fact.predicate, fact.object, fact.raw_text,
+                fact.confidence, fact.status.value,
+                _dt_to_str(fact.valid_from), _dt_to_str(fact.valid_to),
+                fact.superseded_by, fact.id,
+            ),
+        )
+        self._conn.commit()
+        return fact
+
+    # ── Relations ──
+
+    def add_relation(self, relation: Relation) -> Relation:
+        relation.created_at = datetime.now()
+        self._conn.execute(
+            """INSERT INTO relations
+               (id, from_entity_id, to_entity_id, relation_type, metadata,
+                valid_from, valid_to, created_at)
+               VALUES (?,?,?,?,?,?,?,?)""",
+            (
+                relation.id, relation.from_entity_id, relation.to_entity_id,
+                relation.relation_type,
+                json.dumps(relation.metadata, ensure_ascii=False),
+                _dt_to_str(relation.valid_from), _dt_to_str(relation.valid_to),
+                _dt_to_str(relation.created_at),
+            ),
+        )
+        self._conn.commit()
+        return relation
+
+    def get_relations(self, entity_id: str, direction: str = "both") -> list[Relation]:
+        if direction == "outgoing":
+            rows = self._conn.execute(
+                "SELECT * FROM relations WHERE from_entity_id = ?", (entity_id,)
+            ).fetchall()
+        elif direction == "incoming":
+            rows = self._conn.execute(
+                "SELECT * FROM relations WHERE to_entity_id = ?", (entity_id,)
+            ).fetchall()
+        else:
+            rows = self._conn.execute(
+                "SELECT * FROM relations WHERE from_entity_id = ? OR to_entity_id = ?",
+                (entity_id, entity_id),
+            ).fetchall()
+        return [self._row_to_relation(r) for r in rows]
+
+    # ── Sessions ──
+
+    def save_session(self, session: Session) -> None:
+        self._conn.execute(
+            """INSERT OR REPLACE INTO sessions
+               (session_id, agent_id, started_at, ended_at, parent_session_id,
+                entities_accessed, entities_modified, facts_added, summary, metadata)
+               VALUES (?,?,?,?,?,?,?,?,?,?)""",
+            (
+                session.session_id, session.agent_id,
+                _dt_to_str(session.started_at), _dt_to_str(session.ended_at),
+                session.parent_session_id,
+                json.dumps(session.entities_accessed),
+                json.dumps(session.entities_modified),
+                json.dumps(session.facts_added),
+                session.summary,
+                json.dumps(session.metadata, ensure_ascii=False),
+            ),
+        )
+        self._conn.commit()
+
+    def get_session(self, session_id: str) -> Optional[Session]:
+        row = self._conn.execute(
+            "SELECT * FROM sessions WHERE session_id = ?", (session_id,)
+        ).fetchone()
+        return self._row_to_session(row) if row else None
+
+    def get_recent_sessions(self, agent_id: str, limit: int = 5) -> list[Session]:
+        rows = self._conn.execute(
+            "SELECT * FROM sessions WHERE agent_id = ? ORDER BY started_at DESC LIMIT ?",
+            (agent_id, limit),
+        ).fetchall()
+        return [self._row_to_session(r) for r in rows]
+
+    def add_session_event(self, event: SessionEvent) -> None:
+        self._conn.execute(
+            """INSERT INTO session_events (session_id, event_type, target, detail, timestamp)
+               VALUES (?,?,?,?,?)""",
+            (
+                event.session_id, event.event_type, event.target,
+                json.dumps(event.detail, ensure_ascii=False),
+                _dt_to_str(event.timestamp),
+            ),
+        )
+        self._conn.commit()
+
+    # ── Conflicts ──
+
+    def add_conflict(self, conflict: dict) -> None:
+        # Validate required keys
+        if "id" not in conflict or "conflict_type" not in conflict:
+            raise StorageError("Conflict dict must contain 'id' and 'conflict_type' keys")
+        self._conn.execute(
+            """INSERT INTO conflicts
+               (id, existing_fact_id, new_fact_id, conflict_type, suggested_resolution)
+               VALUES (?,?,?,?,?)""",
+            (
+                conflict["id"], conflict.get("existing_fact_id"),
+                conflict.get("new_fact_id"), conflict["conflict_type"],
+                conflict.get("suggested_resolution"),
+            ),
+        )
+        self._conn.commit()
+
+    def get_pending_conflicts(self) -> list[dict]:
+        rows = self._conn.execute(
+            "SELECT * FROM conflicts WHERE resolved_at IS NULL"
+        ).fetchall()
+        return [dict(r) for r in rows]
+
+    def resolve_conflict(self, conflict_id: str, resolution: str, resolved_by: str) -> None:
+        self._conn.execute(
+            """UPDATE conflicts SET resolution=?, resolved_at=?, resolved_by=? WHERE id=?""",
+            (resolution, _dt_to_str(datetime.now()), resolved_by, conflict_id),
+        )
+        self._conn.commit()
+
+    # ── Stats ──
+
+    def entity_count(self) -> int:
+        row = self._conn.execute("SELECT COUNT(*) FROM entities").fetchone()
+        return row[0] if row else 0
+
+    def fact_count(self) -> int:
+        row = self._conn.execute("SELECT COUNT(*) FROM facts").fetchone()
+        return row[0] if row else 0
+
+    # ── FTS Indexing ──
+
+    def _index_entity(self, entity: Entity, facts: Optional[list[Fact]] = None) -> None:
+        if facts is None:
+            facts = self.get_current_facts(entity.id)
+        # Only index top 50 most recent/confident facts to prevent FTS bloat
+        relevant = sorted(facts, key=lambda f: (f.confidence, f.created_at.isoformat()), reverse=True)[:50]
+        fact_texts = " ".join(f.raw_text for f in relevant)
+
+        # Include relation targets so searching "Hashed" also finds "Simon Kim (CEO_OF Hashed)"
+        relations = self.get_relations(entity.id)
+        relation_parts: list[str] = []
+        for r in relations[:20]:
+            other_id = r.to_entity_id if r.from_entity_id == entity.id else r.from_entity_id
+            other = self.get_entity(other_id)
+            if other:
+                relation_parts.append(f"{r.relation_type} {other.name}")
+        if relation_parts:
+            fact_texts = f"{fact_texts} {' '.join(relation_parts)}"
+
+        state_text = " ".join(f"{k}: {v}" for k, v in entity.state.to_dict().items() if isinstance(v, str))
+        # Include aliases in searchable text
+        alias_text = " ".join(entity.aliases) if entity.aliases else ""
+        if alias_text:
+            fact_texts = f"{fact_texts} {alias_text}"
+
+        self._conn.execute(
+            "INSERT INTO memory_fts (entity_id, entity_name, entity_type, fact_text, state_text, summary) VALUES (?,?,?,?,?,?)",
+            (entity.id, entity.name, entity.entity_type, fact_texts, state_text, entity.summary),
+        )
+
+    def _reindex_entity(self, entity: Entity) -> None:
+        self._conn.execute("DELETE FROM memory_fts WHERE entity_id = ?", (entity.id,))
+        facts = self.get_facts(entity.id)
+        self._index_entity(entity, facts)
+
+    def reindex_all(self) -> int:
+        self._conn.execute("DELETE FROM memory_fts")
+        count = 0
+        offset = 0
+        batch_size = 500
+        while True:
+            batch = self.list_entities(limit=batch_size, offset=offset)
+            if not batch:
+                break
+            for entity in batch:
+                self._index_entity(entity)
+            count += len(batch)
+            offset += batch_size
+        self._conn.commit()
+        return count
+
+    # ── Row → Model Conversion ──
+
+    def _row_to_entity(self, row: sqlite3.Row) -> Entity:
+        state_data = json.loads(row["state"]) if row["state"] else {}
+        return Entity(
+            id=row["id"],
+            name=row["name"],
+            entity_type=row["entity_type"],
+            tier=Tier(row["tier"]),
+            summary=row["summary"] or "",
+            aliases=json.loads(row["aliases"]) if row["aliases"] else [],
+            state=EntityState.from_dict(state_data),
+            created_at=_str_to_dt(row["created_at"]) or datetime.now(),
+            updated_at=_str_to_dt(row["updated_at"]) or datetime.now(),
+            access_count=row["access_count"] or 0,
+            last_accessed=_str_to_dt(row["last_accessed"]),
+        )
+
+    def _row_to_fact(self, row: sqlite3.Row) -> Fact:
+        return Fact(
+            id=row["id"],
+            entity_id=row["entity_id"],
+            subject=row["subject"],
+            predicate=row["predicate"],
+            object=row["object"],
+            raw_text=row["raw_text"],
+            source=Source(
+                type=SourceType(row["source_type"]),
+                author=row["source_author"] or "",
+                channel=row["source_channel"] or "",
+                timestamp=_str_to_dt(row["source_timestamp"]) or datetime.now(),
+                confidence=row["source_confidence"] if row["source_confidence"] is not None else 1.0,
+                url=row["source_url"],
+                session_id=row["source_session_id"],
+            ),
+            confidence=row["confidence"],
+            status=FactStatus(row["status"]),
+            valid_from=_str_to_dt(row["valid_from"]),
+            valid_to=_str_to_dt(row["valid_to"]),
+            superseded_by=row["superseded_by"],
+            created_at=_str_to_dt(row["created_at"]) or datetime.now(),
+        )
+
+    def _row_to_relation(self, row: sqlite3.Row) -> Relation:
+        return Relation(
+            id=row["id"],
+            from_entity_id=row["from_entity_id"],
+            to_entity_id=row["to_entity_id"],
+            relation_type=row["relation_type"],
+            metadata=json.loads(row["metadata"]) if row["metadata"] else {},
+            valid_from=_str_to_dt(row["valid_from"]),
+            valid_to=_str_to_dt(row["valid_to"]),
+            created_at=_str_to_dt(row["created_at"]) or datetime.now(),
+        )
+
+    def _row_to_session(self, row: sqlite3.Row) -> Session:
+        return Session(
+            session_id=row["session_id"],
+            agent_id=row["agent_id"],
+            started_at=_str_to_dt(row["started_at"]) or datetime.now(),
+            ended_at=_str_to_dt(row["ended_at"]),
+            parent_session_id=row["parent_session_id"],
+            entities_accessed=json.loads(row["entities_accessed"]) if row["entities_accessed"] else [],
+            entities_modified=json.loads(row["entities_modified"]) if row["entities_modified"] else [],
+            facts_added=json.loads(row["facts_added"]) if row["facts_added"] else [],
+            summary=row["summary"],
+            metadata=json.loads(row["metadata"]) if row["metadata"] else {},
+        )