NexusOS/synapse/memory/store.py

from pathlib import Path
from typing import Any, Dict, List, Optional
from pydantic import BaseModel

from ..nexus_config import DEFAULT_MEMORY_MODEL
import sqlite3
import json
import math
import os
import time


def _cosine(a: List[float], b: List[float]) -> float:
    """Cosine similarity between two equal-length vectors. 0.0 on mismatch."""
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0.0 or nb == 0.0:
        return 0.0
    return dot / (na * nb)

# -----------------------------
# Models
# -----------------------------
class MemoryItem(BaseModel):
    id: str
    section: str = "General"
    text: str
    tags: List[str] = []
    position: int = 0
    project_id: str = ""   # "" = global: injected into every chat

class MessageItem(BaseModel):
    role: str  # "user" or "assistant"
    content: str
    timestamp: float
    model: Optional[str] = None
    tokens: Optional[int] = None

class ConversationItem(BaseModel):
    id: str
    messages: List[MessageItem] = []
    created_at: float
    updated_at: float
    title: Optional[str] = None
    project_id: str = ""

    @property
    def preview(self) -> str:
        for msg in self.messages:
            if msg.role == "user":
                return msg.content[:80]
        return "Empty conversation"

    @property
    def timestamp(self) -> float:
        return self.created_at

# -----------------------------
# Persistent Store
# -----------------------------
class PersistentMemoryStore:
    def __init__(self, db_path: Path):
        self.db_path = db_path
        os.makedirs(self.db_path.parent, exist_ok=True)
        self.vec_enabled = self._probe_vec()
        self._ensure_tables()
        self._cache: Dict[str, MemoryItem] = self._load_all_memory()

    # -----------------------------
    # Internal helpers
    # -----------------------------
    @staticmethod
    def _probe_vec() -> bool:
        """True if this host can load the sqlite-vec extension. Some Python
        builds ship SQLite with loadable extensions disabled — those fall back
        to the brute-force cosine scan, so recall never depends on this."""
        try:
            import sqlite_vec
            c = sqlite3.connect(":memory:")
            c.enable_load_extension(True)
            sqlite_vec.load(c)
            c.execute("SELECT vec_version()")
            c.close()
            return True
        except Exception:
            return False

    def _connect(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL;")
        if self.vec_enabled:
            try:
                import sqlite_vec
                conn.enable_load_extension(True)
                sqlite_vec.load(conn)
                conn.enable_load_extension(False)
            except Exception:
                pass
        return conn

    def _ensure_tables(self):
        conn = self._connect()
        cur = conn.cursor()

        cur.execute("""
            CREATE TABLE IF NOT EXISTS memory (
                id TEXT PRIMARY KEY,
                section TEXT NOT NULL DEFAULT 'General',
                text TEXT NOT NULL,
                tags TEXT
            )
        """)
        # Migrate: add section column if it doesn't exist yet
        try:
            cur.execute("ALTER TABLE memory ADD COLUMN section TEXT NOT NULL DEFAULT 'General'")
        except Exception:
            pass
        # Migrate: add position column for stable ordering
        try:
            cur.execute("ALTER TABLE memory ADD COLUMN position INTEGER NOT NULL DEFAULT 0")
        except Exception:
            pass
        # Backfill positions for rows added before this column existed
        cur.execute("SELECT COUNT(*) FROM memory WHERE position > 0")
        if cur.fetchone()[0] == 0:
            cur.execute("UPDATE memory SET position = rowid")

        # Migrate: scope a fact to a project ("" = global, applies everywhere).
        try:
            cur.execute("ALTER TABLE memory ADD COLUMN project_id TEXT NOT NULL DEFAULT ''")
        except Exception:
            pass

        cur.execute("""
            CREATE TABLE IF NOT EXISTS conversations (
                id TEXT PRIMARY KEY,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL,
                title TEXT
            )
        """)
        # Migrate: add title column if it doesn't exist yet
        try:
            cur.execute("ALTER TABLE conversations ADD COLUMN title TEXT")
        except Exception:
            pass
        # Migrate: bind a conversation to a project ("" = unscoped).
        try:
            cur.execute("ALTER TABLE conversations ADD COLUMN project_id TEXT NOT NULL DEFAULT ''")
        except Exception:
            pass
        # Migrate: high-water mark for memory extraction — the id of the last
        # message the curator has already read. Extraction runs once the
        # conversation goes idle rather than after every exchange, so this is
        # what makes it idempotent and restart-safe: a backend that dies with a
        # pending sweep resumes from here instead of re-reading the whole
        # transcript and re-saving facts it already saved.
        try:
            cur.execute("ALTER TABLE conversations ADD COLUMN extracted_through INTEGER NOT NULL DEFAULT 0")
            # Only reached the first time the column is added. Everything already
            # in the database was extracted per-exchange under the old design, so
            # watermark it as read — without this the first idle sweep would
            # re-read every historical conversation and re-save its facts.
            cur.execute(
                "UPDATE conversations SET extracted_through = "
                "COALESCE((SELECT MAX(id) FROM messages WHERE conversation_id = conversations.id), 0)"
            )
        except Exception:
            pass

        cur.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp REAL NOT NULL,
                model TEXT,
                tokens INTEGER,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)
        for col, typedef in (("model", "TEXT"), ("tokens", "INTEGER")):
            try:
                cur.execute(f"ALTER TABLE messages ADD COLUMN {col} {typedef}")
            except Exception:
                pass

        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_messages_conversation_id
            ON messages (conversation_id)
        """)

        # Semantic recall: one embedding vector per message, stored as JSON.
        # Backfilled lazily by semantic_search_conversations so existing history
        # gets indexed on first search.
        cur.execute("""
            CREATE TABLE IF NOT EXISTS message_vectors (
                message_id INTEGER PRIMARY KEY,
                embedding TEXT NOT NULL,
                FOREIGN KEY (message_id) REFERENCES messages(id)
            )
        """)

        # RAG: one row per chunk. doc_id groups the chunks of a single uploaded
        # document; embedding is a JSON vector (search_document-prefixed), set at
        # ingest so retrieval needs no backfill.
        cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id TEXT PRIMARY KEY,
                doc_id TEXT NOT NULL,
                title TEXT NOT NULL,
                chunk_idx INTEGER NOT NULL,
                text TEXT NOT NULL,
                embedding TEXT,
                created_at REAL NOT NULL
            )
        """)
        cur.execute("CREATE INDEX IF NOT EXISTS idx_documents_doc_id ON documents (doc_id)")

        # Projects / workspaces: group documents so RAG can scope to one set.
        cur.execute("""
            CREATE TABLE IF NOT EXISTS projects (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                created_at REAL NOT NULL
            )
        """)
        # projects.instructions — per-project system prompt ("" = none).
        try:
            cur.execute("ALTER TABLE projects ADD COLUMN instructions TEXT NOT NULL DEFAULT ''")
        except Exception:
            pass
        # documents.project_id — "" (or missing) means unscoped / All.
        try:
            cur.execute("ALTER TABLE documents ADD COLUMN project_id TEXT NOT NULL DEFAULT ''")
        except Exception:
            pass

        cur.execute("""
            CREATE TABLE IF NOT EXISTS settings (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )
        """)
        cur.execute(
            "DELETE FROM settings WHERE key IN ('anthropic_api_key', 'escalation_model')"
        )

        self._sweep_orphan_msg_vectors(conn)

        conn.commit()
        conn.close()

    # -----------------------------
    # Loaders
    # -----------------------------
    def _load_all_memory(self) -> Dict[str, MemoryItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT id, section, text, tags, position, project_id FROM memory ORDER BY position ASC, rowid ASC")
        rows = cur.fetchall()
        conn.close()

        cache = {}
        for row in rows:
            try:
                tags = json.loads(row["tags"]) if row["tags"] else []
            except Exception:
                tags = []
            cache[row["id"]] = MemoryItem(
                id=row["id"],
                section=row["section"] or "General",
                text=row["text"],
                tags=tags,
                position=row["position"] or 0,
                project_id=row["project_id"] or "",
            )

        return cache

    # -----------------------------
    # Memory API
    # -----------------------------
    def add(self, item: MemoryItem):
        if not item.position:
            conn = self._connect()
            try:
                row = conn.execute(
                    "SELECT position FROM memory WHERE id = ?", (item.id,)
                ).fetchone()
                if row and row["position"]:
                    item.position = row["position"]
                else:
                    row = conn.execute("SELECT MAX(position) AS max_position FROM memory").fetchone()
                    item.position = (row["max_position"] or 0) + 1
            finally:
                conn.close()
        self._cache[item.id] = item
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT OR REPLACE INTO memory (id, section, text, tags, position, project_id)"
                " VALUES (?, ?, ?, ?, ?, ?)",
                (item.id, item.section or "General", item.text, json.dumps(item.tags),
                 item.position, item.project_id or "")
            )
            conn.commit()
        finally:
            conn.close()

    def update(self, item: MemoryItem):
        existing = self.get(item.id)
        if existing:
            item.position = existing.position
        self.add(item)

    def delete(self, item_id: str):
        self._cache.pop(item_id, None)
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute("DELETE FROM memory WHERE id = ?", (item_id,))
            conn.commit()
        finally:
            conn.close()

    def get(self, item_id: str) -> Optional[MemoryItem]:
        conn = self._connect()
        try:
            row = conn.execute(
                "SELECT id, section, text, tags, position FROM memory WHERE id = ?",
                (item_id,),
            ).fetchone()
        finally:
            conn.close()
        if not row:
            self._cache.pop(item_id, None)
            return None
        try:
            tags = json.loads(row["tags"]) if row["tags"] else []
        except Exception:
            tags = []
        item = MemoryItem(
            id=row["id"], section=row["section"] or "General", text=row["text"],
            tags=tags, position=row["position"] or 0,
        )
        self._cache[item_id] = item
        return item

    def all(self) -> List[MemoryItem]:
        # Always read from DB, never the cache: the CLI, the control panel and
        # the backend are separate processes with separate caches, so a fact
        # written by one is invisible to another's cache.
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT id, section, text, tags, position, project_id FROM memory ORDER BY position ASC, rowid ASC")
        rows = cur.fetchall()
        conn.close()
        items = []
        for row in rows:
            try:
                tags = json.loads(row["tags"]) if row["tags"] else []
            except Exception:
                tags = []
            items.append(MemoryItem(
                id=row["id"],
                section=row["section"] or "General",
                text=row["text"],
                tags=tags,
                position=row["position"] or 0,
                project_id=row["project_id"] or "",
            ))
        return items

    def reorder_section(self, section: str, ordered_ids: List[str]) -> bool:
        """Rewrite the order of items in a section using its existing position pool.
        ordered_ids must contain exactly the ids currently in the section."""
        section_norm = section or "General"
        in_section = [i for i in self.all() if (i.section or "General") == section_norm]
        if len(in_section) != len(ordered_ids):
            return False
        by_id = {i.id: i for i in in_section}
        siblings = []
        for id_ in ordered_ids:
            if id_ not in by_id:
                return False
            siblings.append(by_id[id_])
        positions = sorted([i.position for i in in_section])
        conn = self._connect()
        try:
            cur = conn.cursor()
            for s, new_pos in zip(siblings, positions):
                s.position = new_pos
                cur.execute("UPDATE memory SET position = ? WHERE id = ?", (new_pos, s.id))
            conn.commit()
        finally:
            conn.close()
        return True

    # -----------------------------
    # Conversation API
    # -----------------------------
    def create_conversation(self, conversation_id: str, project_id: str = "") -> ConversationItem:
        now = time.time()
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT OR IGNORE INTO conversations (id, created_at, updated_at, project_id) VALUES (?, ?, ?, ?)",
                (conversation_id, now, now, project_id or "")
            )
            conn.commit()
        finally:
            conn.close()
        return ConversationItem(id=conversation_id, created_at=now, updated_at=now)

    def conversation_project(self, conversation_id: str) -> Optional[str]:
        """The conversation's bound project ('' = unscoped), or None if it doesn't
        exist yet — lets a new chat inherit the current workspace."""
        conn = self._connect()
        row = conn.execute(
            "SELECT project_id FROM conversations WHERE id = ?", (conversation_id,)
        ).fetchone()
        conn.close()
        return None if row is None else (row["project_id"] or "")

    def set_conversation_title(self, conversation_id: str, title: str):
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "UPDATE conversations SET title = ? WHERE id = ?",
                (title, conversation_id),
            )
            conn.commit()
        finally:
            conn.close()

    def set_conversation_project(self, conversation_id: str, project_id: str):
        """Move a conversation into a project ('' = unscoped)."""
        conn = self._connect()
        try:
            conn.execute(
                "UPDATE conversations SET project_id = ? WHERE id = ?",
                (project_id or "", conversation_id),
            )
            conn.commit()
        finally:
            conn.close()

    def add_message(self, conversation_id: str, role: str, content: str, model: Optional[str] = None, tokens: Optional[int] = None):
        now = time.time()
        conn = self._connect()
        try:
            cur = conn.cursor()
            cur.execute(
                "INSERT INTO messages (conversation_id, role, content, timestamp, model, tokens) VALUES (?, ?, ?, ?, ?, ?)",
                (conversation_id, role, content, now, model, tokens)
            )
            message_id = cur.lastrowid
            cur.execute(
                "UPDATE conversations SET updated_at = ? WHERE id = ?",
                (now, conversation_id)
            )
            conn.commit()
            return message_id
        finally:
            conn.close()

    def get_conversation(self, conversation_id: str) -> Optional[ConversationItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT * FROM conversations WHERE id = ?", (conversation_id,))
        row = cur.fetchone()
        if not row:
            conn.close()
            return None
        cur.execute(
            "SELECT role, content, timestamp, model, tokens FROM messages WHERE conversation_id = ? ORDER BY timestamp ASC",
            (conversation_id,)
        )
        msg_rows = cur.fetchall()
        conn.close()

        messages = [MessageItem(role=r["role"], content=r["content"], timestamp=r["timestamp"], model=r["model"], tokens=r["tokens"]) for r in msg_rows]
        return ConversationItem(
            id=row["id"],
            messages=messages,
            created_at=row["created_at"],
            updated_at=row["updated_at"],
            title=row["title"] if "title" in row.keys() else None,
        )

    def all_conversations(self) -> List[ConversationItem]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("""
            SELECT c.id, c.created_at, c.updated_at, c.title, c.project_id,
                   m.role, m.content, m.timestamp, m.model, m.tokens
            FROM conversations c
            LEFT JOIN messages m ON m.conversation_id = c.id
            ORDER BY c.updated_at DESC, m.timestamp ASC
        """)
        rows = cur.fetchall()
        conn.close()

        convs: Dict[str, ConversationItem] = {}
        order: list[str] = []
        for row in rows:
            cid = row["id"]
            if cid not in convs:
                convs[cid] = ConversationItem(
                    id=cid,
                    created_at=row["created_at"],
                    updated_at=row["updated_at"],
                    title=row["title"],
                    project_id=row["project_id"] or "",
                )
                order.append(cid)
            if row["role"] is not None:
                convs[cid].messages.append(
                    MessageItem(role=row["role"], content=row["content"], timestamp=row["timestamp"], model=row["model"], tokens=row["tokens"])
                )
        return [convs[cid] for cid in order]

    def pending_extraction(self, conversation_id: str) -> tuple[list, int]:
        """Messages the curator has not read yet, and the id to watermark to.

        Returns ([{role, content}], last_id). An empty list means nothing new,
        so callers can skip the model call entirely.
        """
        conn = self._connect()
        try:
            row = conn.execute(
                "SELECT extracted_through FROM conversations WHERE id = ?", (conversation_id,)
            ).fetchone()
            if row is None:
                return [], 0
            rows = conn.execute(
                "SELECT id, role, content FROM messages "
                "WHERE conversation_id = ? AND id > ? ORDER BY id ASC",
                (conversation_id, row["extracted_through"] or 0),
            ).fetchall()
        finally:
            conn.close()
        if not rows:
            return [], 0
        return ([{"role": r["role"], "content": r["content"]} for r in rows], rows[-1]["id"])

    def set_extracted_through(self, conversation_id: str, message_id: int) -> None:
        conn = self._connect()
        try:
            conn.execute(
                "UPDATE conversations SET extracted_through = ? WHERE id = ?",
                (int(message_id), conversation_id),
            )
            conn.commit()
        finally:
            conn.close()

    def conversations_awaiting_extraction(self, idle_seconds: float) -> list[str]:
        """Conversations with unread messages that have been quiet long enough
        to count as finished. Used to resume sweeps dropped by a restart."""
        conn = self._connect()
        try:
            rows = conn.execute(
                "SELECT c.id FROM conversations c JOIN messages m ON m.conversation_id = c.id "
                "WHERE m.id > c.extracted_through GROUP BY c.id "
                "HAVING MAX(m.timestamp) < ?",
                (time.time() - idle_seconds,),
            ).fetchall()
        finally:
            conn.close()
        return [r["id"] for r in rows]

    def delete_conversation(self, conversation_id: str):
        conn = self._connect()
        try:
            cur = conn.cursor()
            # Drop the embeddings first, while the message ids still resolve -
            # the delete-side counterpart of _vec_upsert_msg (see its docstring).
            ids = [r["id"] for r in cur.execute(
                "SELECT id FROM messages WHERE conversation_id = ?", (conversation_id,)
            ).fetchall()]
            self._delete_message_vectors(conn, ids)
            cur.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
            cur.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
            conn.commit()
        finally:
            conn.close()

    # -----------------------------
    # Search API
    # -----------------------------
    def search_conversations(self, query: str, limit: int = 3) -> List[dict]:
        """Return up to `limit` conversations that contain the query string,
        with full user+assistant exchange pairs around each match."""
        if not query or not query.strip():
            return []
        q = query.strip().lower()
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("""
            SELECT DISTINCT c.id, c.created_at, c.updated_at
            FROM conversations c
            JOIN messages m ON m.conversation_id = c.id
            WHERE LOWER(m.content) LIKE ?
            ORDER BY c.updated_at DESC
            LIMIT ?
        """, (f"%{q}%", limit))
        rows = cur.fetchall()

        results = []
        for row in rows:
            # Load all messages in order so we can find complete exchange pairs
            cur.execute("""
                SELECT role, content FROM messages
                WHERE conversation_id = ?
                ORDER BY timestamp ASC
            """, (row["id"],))
            all_msgs = [{"role": r["role"], "content": r["content"]} for r in cur.fetchall()]

            # For each matching message, collect the full user+assistant pair around it
            seen_pairs: set = set()
            matches = []
            for i, msg in enumerate(all_msgs):
                if q not in msg["content"].lower():
                    continue
                if msg["role"] == "user":
                    start, end = i, i + 1 if i + 1 < len(all_msgs) else i
                else:
                    start, end = (i - 1 if i > 0 else i), i
                if (start, end) in seen_pairs:
                    continue
                seen_pairs.add((start, end))
                for m in all_msgs[start:end + 1]:
                    matches.append({"role": m["role"], "content": m["content"][:500]})
                if len(seen_pairs) >= 2:
                    break

            if matches:
                results.append({
                    "id": row["id"],
                    "updated_at": row["updated_at"],
                    "matches": matches,
                })

        conn.close()
        return results

    # nomic-embed-text is an asymmetric retrieval model: queries and stored
    # documents must be embedded with these task prefixes or similarity collapses
    # into noise. Cached vectors are document-embeddings (search_document:).
    _EMBED_QUERY_PREFIX = "search_query: "
    _EMBED_DOC_PREFIX = "search_document: "

    async def semantic_search_conversations(
        self, query: str, embed_fn, limit: int = 3, min_score: float = 0.6
    ) -> List[dict]:
        """Recall past conversations relevant to `query` using hybrid retrieval.

        Combines semantic similarity (embeddings — finds reworded matches with no
        shared keywords) with the existing lexical substring match (catches exact
        terms the embedding underweights), unioned and deduped by conversation.

        `embed_fn` is an async callable returning an embedding vector for a string
        (typically OllamaManager.embed). Messages without a stored vector are
        embedded and cached on first use (lazy backfill). If embeddings are
        unavailable (no model / Ollama down) this degrades to pure lexical match,
        so recall never silently breaks.

        Returns the same shape as `search_conversations`: a list of
        {id, updated_at, matches:[{role, content}]} with full user+assistant
        pairs around each match.
        """
        if not query or not query.strip():
            return []

        query_vec = await embed_fn(self._EMBED_QUERY_PREFIX + query.strip())
        if not query_vec:
            return self.search_conversations(query, limit=limit)

        conn = self._connect()
        cur = conn.cursor()

        # Lazy backfill: embed any messages that don't have a vector yet.
        cur.execute("""
            SELECT m.id, m.content
            FROM messages m
            LEFT JOIN message_vectors v ON v.message_id = m.id
            WHERE v.message_id IS NULL AND TRIM(m.content) != ''
        """)
        missing = cur.fetchall()
        for row in missing:
            vec = await embed_fn(self._EMBED_DOC_PREFIX + row["content"][:2000])
            if vec:
                cur.execute(
                    "INSERT OR REPLACE INTO message_vectors (message_id, embedding) VALUES (?, ?)",
                    (row["id"], json.dumps(vec)),
                )
                self._vec_upsert_msg(conn, row["id"], vec)  # mirror into the ANN index
        if missing:
            conn.commit()

        # Rank messages by similarity. Fast path: the sqlite-vec ANN index over an
        # over-fetch (conversation dedup below thins it); else brute-force cosine.
        scored = None
        if self.vec_enabled:
            scored = self._vec_search_messages(conn, query_vec, max(limit * 5, 20), min_score)
        if scored is None:
            cur.execute("""
                SELECT v.message_id, v.embedding, m.conversation_id
                FROM message_vectors v
                JOIN messages m ON m.id = v.message_id
            """)
            scored = []
            for row in cur.fetchall():
                try:
                    vec = json.loads(row["embedding"])
                except Exception:
                    continue
                score = _cosine(query_vec, vec)
                if score >= min_score:
                    scored.append((score, row["message_id"], row["conversation_id"]))
            scored.sort(reverse=True)

        results: List[dict] = []
        seen_convs: set = set()
        for score, message_id, conv_id in scored:
            if len(results) >= limit:
                break
            if conv_id in seen_convs:
                continue
            pair = self._exchange_pair(cur, conv_id, message_id)
            if pair:
                seen_convs.add(conv_id)
                results.append(pair)

        conn.close()

        # Hybrid union: fill any remaining slots with lexical matches the
        # embedding missed (e.g. exact proper nouns), skipping dupes.
        if len(results) < limit:
            for conv in self.search_conversations(query, limit=limit):
                if conv["id"] not in seen_convs:
                    seen_convs.add(conv["id"])
                    results.append(conv)
                    if len(results) >= limit:
                        break

        return results

    # --- message vector index (sqlite-vec) — same pattern as documents --------
    def _ensure_vec_msgs(self, conn, dim: int) -> None:
        conn.execute(
            f"CREATE VIRTUAL TABLE IF NOT EXISTS vec_messages "
            f"USING vec0(embedding float[{dim}] distance_metric=cosine)"
        )

    def _vec_upsert_msg(self, conn, message_id: int, vec: list) -> None:
        if not (self.vec_enabled and vec):
            return
        try:
            import sqlite_vec
            self._ensure_vec_msgs(conn, len(vec))
            conn.execute("DELETE FROM vec_messages WHERE rowid = ?", (message_id,))
            conn.execute(
                "INSERT INTO vec_messages(rowid, embedding) VALUES (?, ?)",
                (message_id, sqlite_vec.serialize_float32(vec)),
            )
        except Exception:
            pass

    def _delete_message_vectors(self, conn, message_ids) -> None:
        """Drop the cached embeddings of messages that are about to be deleted,
        mirroring the removal into the ANN index — the delete-side counterpart of
        `_vec_upsert_msg`. Retrieval already ignores orphans (it inner-joins
        messages), but `messages.id` is AUTOINCREMENT so a stale vector is never
        overwritten either: without this the table and index only ever grow."""
        ids = list(message_ids)
        if not ids:
            return
        for i in range(0, len(ids), 500):   # stay under SQLite's variable limit
            batch = ids[i:i + 500]
            conn.execute(
                f"DELETE FROM message_vectors WHERE message_id IN ({','.join('?' * len(batch))})",
                batch,
            )
        if self.vec_enabled:
            try:
                for mid in ids:
                    conn.execute("DELETE FROM vec_messages WHERE rowid = ?", (mid,))
            except Exception:
                pass  # index absent / extension unavailable — it's only a mirror

    def _sweep_orphan_msg_vectors(self, conn) -> None:
        """One-time repair for databases written before delete_conversation
        cleaned up after itself: drop vectors whose message is already gone."""
        try:
            ids = [r["message_id"] for r in conn.execute(
                "SELECT v.message_id FROM message_vectors v "
                "LEFT JOIN messages m ON m.id = v.message_id WHERE m.id IS NULL"
            ).fetchall()]
            if ids:
                self._delete_message_vectors(conn, ids)
        except Exception:
            pass

    def _backfill_vec_msgs(self, conn, dim: int) -> None:
        """Index any message_vectors rows missing from vec_messages."""
        try:
            import sqlite_vec
            self._ensure_vec_msgs(conn, dim)
            rows = conn.execute(
                "SELECT mv.message_id AS mid, mv.embedding AS emb FROM message_vectors mv "
                "LEFT JOIN vec_messages v ON v.rowid = mv.message_id WHERE v.rowid IS NULL"
            ).fetchall()
            for r in rows:
                try:
                    vec = json.loads(r["emb"])
                    if len(vec) == dim:
                        conn.execute(
                            "INSERT INTO vec_messages(rowid, embedding) VALUES (?, ?)",
                            (r["mid"], sqlite_vec.serialize_float32(vec)),
                        )
                except Exception:
                    pass
            conn.commit()
        except Exception:
            pass

    def _vec_search_messages(self, conn, query_vec: list, fetch: int, min_score: float):
        """Top message hits via the vec index as sorted [(score, msg_id, conv_id)],
        or None to fall back to the brute-force scan. Stale rows for deleted
        messages are dropped by the inner join, so they never surface."""
        try:
            import sqlite_vec
            self._backfill_vec_msgs(conn, len(query_vec))
            rows = conn.execute(
                "SELECT v.rowid AS mid, v.distance AS distance, m.conversation_id AS cid "
                "FROM vec_messages v JOIN messages m ON m.id = v.rowid "
                "WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
                (sqlite_vec.serialize_float32(query_vec), fetch),
            ).fetchall()
            return [
                (1.0 - r["distance"], r["mid"], r["cid"])
                for r in rows if (1.0 - r["distance"]) >= min_score
            ]  # distance-asc == score-desc, already sorted
        except Exception:
            return None

    def _exchange_pair(self, cur, conv_id: str, message_id: int) -> Optional[dict]:
        """Build a {id, updated_at, matches} record with the full user+assistant
        pair surrounding `message_id`, in the shape search callers expect."""
        cur.execute(
            "SELECT id, role, content FROM messages WHERE conversation_id = ? ORDER BY timestamp ASC",
            (conv_id,),
        )
        all_msgs = cur.fetchall()
        idx = next((i for i, m in enumerate(all_msgs) if m["id"] == message_id), None)
        if idx is None:
            return None
        if all_msgs[idx]["role"] == "user":
            start, end = idx, min(idx + 1, len(all_msgs) - 1)
        else:
            start, end = max(idx - 1, 0), idx
        matches = [
            {"role": all_msgs[i]["role"], "content": all_msgs[i]["content"][:500]}
            for i in range(start, end + 1)
        ]
        cur.execute("SELECT updated_at FROM conversations WHERE id = ?", (conv_id,))
        crow = cur.fetchone()
        return {
            "id": conv_id,
            "updated_at": crow["updated_at"] if crow else 0,
            "matches": matches,
        }

    # -----------------------------
    # Documents (RAG)
    # -----------------------------
    @staticmethod
    def _chunk_text(text: str, size: int = 800, overlap: int = 120) -> List[str]:
        """Pack paragraphs into ~`size`-char chunks with a char `overlap` carried
        across boundaries, so a passage spanning two chunks still matches. Any
        single paragraph larger than `size` (common in PDFs with few blank lines)
        is hard-split into overlapping windows first.
        ponytail: char-based, not token-based — fine for local recall; move to a
        token splitter only if chunk sizes start hurting the context budget."""
        units: List[str] = []
        for para in (p.strip() for p in text.split("\n\n")):
            if not para:
                continue
            if len(para) <= size:
                units.append(para)
            else:
                step = max(1, size - overlap)
                units.extend(para[i:i + size] for i in range(0, len(para), step))

        chunks: List[str] = []
        buf = ""
        for u in units:
            if buf and len(buf) + len(u) + 2 > size:
                chunks.append(buf)
                tail = buf[-overlap:] if overlap else ""       # overlap seed for the next chunk
                buf = f"{tail}\n\n{u}" if tail else u
            else:
                buf = f"{buf}\n\n{u}" if buf else u
        if buf:
            chunks.append(buf)
        return chunks

    async def add_document(self, title: str, content: str, embed_fn, project_id: str = "") -> dict:
        """Chunk, embed, and store a document. Returns {doc_id, chunks}."""
        import uuid as _uuid
        doc_id = str(_uuid.uuid4())
        pieces = self._chunk_text(content)
        now = time.time()
        conn = self._connect()
        cur = conn.cursor()
        for i, piece in enumerate(pieces):
            vec = await embed_fn(self._EMBED_DOC_PREFIX + piece[:2000])
            cur.execute(
                "INSERT INTO documents (id, doc_id, title, chunk_idx, text, embedding, created_at, project_id)"
                " VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (str(_uuid.uuid4()), doc_id, title, i, piece,
                 json.dumps(vec) if vec else None, now, project_id or ""),
            )
            self._vec_upsert(conn, cur.lastrowid, vec)  # mirror into the ANN index
        conn.commit()
        conn.close()
        return {"doc_id": doc_id, "title": title, "chunks": len(pieces)}

    # --- projects / workspaces --------------------------------------------------
    def create_project(self, name: str) -> dict:
        import uuid as _uuid
        pid = str(_uuid.uuid4())
        conn = self._connect()
        conn.execute("INSERT INTO projects (id, name, created_at) VALUES (?, ?, ?)",
                     (pid, name.strip(), time.time()))
        conn.commit()
        conn.close()
        return {"id": pid, "name": name.strip()}

    def list_projects(self) -> List[dict]:
        conn = self._connect()
        rows = conn.execute("""
            SELECT p.id, p.name, p.created_at, p.instructions,
                   (SELECT COUNT(DISTINCT doc_id) FROM documents d WHERE d.project_id = p.id) AS docs,
                   (SELECT COUNT(*) FROM conversations c WHERE c.project_id = p.id) AS chats
            FROM projects p ORDER BY p.created_at ASC
        """).fetchall()
        conn.close()
        return [dict(r) for r in rows]

    def set_project_instructions(self, project_id: str, instructions: str) -> bool:
        """Per-project system prompt, layered into chats bound to the project."""
        conn = self._connect()
        try:
            cur = conn.execute("UPDATE projects SET instructions = ? WHERE id = ?",
                               (instructions or "", project_id))
            conn.commit()
            return cur.rowcount > 0
        finally:
            conn.close()

    def project_instructions(self, project_id: str) -> str:
        """'' when the project has none, or doesn't exist."""
        if not project_id:
            return ""
        conn = self._connect()
        row = conn.execute("SELECT instructions FROM projects WHERE id = ?", (project_id,)).fetchone()
        conn.close()
        return (row["instructions"] or "") if row else ""

    def delete_project(self, project_id: str) -> bool:
        """Delete a project; its documents, chats and facts survive but become
        unscoped ("")."""
        conn = self._connect()
        for table in ("documents", "conversations", "memory"):
            conn.execute(f"UPDATE {table} SET project_id = '' WHERE project_id = ?", (project_id,))
        cur = conn.execute("DELETE FROM projects WHERE id = ?", (project_id,))
        deleted = cur.rowcount
        conn.commit()
        conn.close()
        return deleted > 0

    # --- vector index (sqlite-vec) — accelerator over the JSON embedding column,
    # --- with brute-force cosine below as the guaranteed fallback ---------------
    def _ensure_vec_docs(self, conn, dim: int) -> None:
        conn.execute(
            f"CREATE VIRTUAL TABLE IF NOT EXISTS vec_documents "
            f"USING vec0(embedding float[{dim}] distance_metric=cosine)"
        )

    def _vec_upsert(self, conn, rowid: int, vec: list) -> None:
        """Best-effort mirror of one chunk's vector into the vec index."""
        if not (self.vec_enabled and vec):
            return
        try:
            import sqlite_vec
            self._ensure_vec_docs(conn, len(vec))
            conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (rowid,))
            conn.execute(
                "INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
                (rowid, sqlite_vec.serialize_float32(vec)),
            )
        except Exception:
            pass  # the index is an accelerator, never a requirement

    def _backfill_vec(self, conn, dim: int) -> None:
        """Index any document chunks missing from vec_documents (older rows, or
        rows written while the extension was unavailable)."""
        try:
            import sqlite_vec
            self._ensure_vec_docs(conn, dim)
            rows = conn.execute(
                "SELECT d.rowid AS rid, d.embedding AS emb FROM documents d "
                "LEFT JOIN vec_documents v ON v.rowid = d.rowid "
                "WHERE v.rowid IS NULL AND d.embedding IS NOT NULL"
            ).fetchall()
            for r in rows:
                try:
                    vec = json.loads(r["emb"])
                    if len(vec) == dim:
                        conn.execute(
                            "INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
                            (r["rid"], sqlite_vec.serialize_float32(vec)),
                        )
                except Exception:
                    pass
            conn.commit()
        except Exception:
            pass

    def _vec_search(self, query_vec: list, limit: int, min_score: float):
        """KNN over the vec index. Returns hits, or None to signal fall back to
        the brute-force scan (e.g. extension error or dimension mismatch)."""
        try:
            import sqlite_vec
            conn = self._connect()
            self._backfill_vec(conn, len(query_vec))
            rows = conn.execute(
                "SELECT d.title AS title, d.text AS text, v.distance AS distance "
                "FROM vec_documents v JOIN documents d ON d.rowid = v.rowid "
                "WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
                (sqlite_vec.serialize_float32(query_vec), limit),
            ).fetchall()
            conn.close()
            # sqlite-vec cosine distance = 1 - cosine similarity
            return [
                {"title": r["title"], "text": r["text"], "score": 1.0 - r["distance"]}
                for r in rows if (1.0 - r["distance"]) >= min_score
            ]
        except Exception:
            return None

    def list_documents(self, project_id: Optional[str] = None) -> List[dict]:
        """All documents, or just one project's when project_id is given."""
        conn = self._connect()
        cur = conn.cursor()
        if project_id is not None:
            cur.execute("""
                SELECT doc_id, title, COUNT(*) AS chunks, MIN(created_at) AS created_at
                FROM documents WHERE project_id = ?
                GROUP BY doc_id, title ORDER BY created_at DESC
            """, (project_id,))
        else:
            cur.execute("""
                SELECT doc_id, title, COUNT(*) AS chunks, MIN(created_at) AS created_at
                FROM documents GROUP BY doc_id, title ORDER BY created_at DESC
            """)
        rows = cur.fetchall()
        conn.close()
        return [dict(r) for r in rows]

    def get_document(self, doc_id: str) -> List[dict]:
        """Ordered chunks of one document: [{chunk_idx, text}]."""
        conn = self._connect()
        cur = conn.cursor()
        cur.execute(
            "SELECT chunk_idx, text FROM documents WHERE doc_id = ? ORDER BY chunk_idx",
            (doc_id,),
        )
        rows = cur.fetchall()
        conn.close()
        return [dict(r) for r in rows]

    def delete_document(self, doc_id: str) -> bool:
        conn = self._connect()
        cur = conn.cursor()
        if self.vec_enabled:
            try:
                for r in cur.execute("SELECT rowid FROM documents WHERE doc_id = ?", (doc_id,)).fetchall():
                    conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (r["rowid"],))
            except Exception:
                pass
        cur.execute("DELETE FROM documents WHERE doc_id = ?", (doc_id,))
        deleted = cur.rowcount
        conn.commit()
        conn.close()
        return deleted > 0

    async def search_documents(
        self, query: str, embed_fn, limit: int = 3, min_score: float = 0.6,
        project_id: Optional[str] = None,
    ) -> List[dict]:
        """Top-`limit` document chunks most similar to `query`. Returns
        [{title, text, score}]. Empty on no query / embeddings down.
        When project_id is given, only that project's docs are searched.
        ponytail: scoped search uses the brute-force path (easy SQL filter, few
        docs per project); the vec index accelerates the unscoped "All" case."""
        if not query or not query.strip():
            return []
        query_vec = await embed_fn(self._EMBED_QUERY_PREFIX + query.strip())
        if not query_vec:
            return []
        # Fast path (unscoped only): the sqlite-vec ANN index.
        if not project_id and self.vec_enabled:
            hits = self._vec_search(query_vec, limit, min_score)
            if hits is not None:
                return hits
        conn = self._connect()
        cur = conn.cursor()
        if project_id:
            cur.execute("SELECT title, text, embedding FROM documents WHERE embedding IS NOT NULL AND project_id = ?", (project_id,))
        else:
            cur.execute("SELECT title, text, embedding FROM documents WHERE embedding IS NOT NULL")
        scored = []
        for row in cur.fetchall():
            try:
                vec = json.loads(row["embedding"])
            except Exception:
                continue
            score = _cosine(query_vec, vec)
            if score >= min_score:
                scored.append({"title": row["title"], "text": row["text"], "score": score})
        conn.close()
        scored.sort(key=lambda d: d["score"], reverse=True)
        return scored[:limit]

    # -----------------------------
    # Settings API
    # -----------------------------
    _SETTINGS_DEFAULTS: Dict[str, Any] = {
        "model": "",
        # Auto-mode routing overrides (blank = built-in preference). When no model
        # is pinned, a message's detected intent picks which of these fires first.
        "auto_chat_model": "",
        "auto_code_model": "",
        # Qwen3-style reasoning. Off by default: the hidden <think> block is pure
        # latency for chat/memory. Turn on for hard multi-step problems.
        "think": False,
        "temperature": 0.7,
        # Context window (tokens Ollama keeps in view). 0 → Ollama's model default.
        "num_ctx": 0,
        # RAG retrieval: how many document chunks to inject, and the minimum
        # cosine similarity (0-1) a chunk must clear to count as relevant.
        "rag_top_k": 3,
        "rag_min_score": 0.6,
        # Active project/workspace; "" = all documents (unscoped).
        "active_project": "",
        # Consent policy for tools that act (web_search/fetch_url/remember):
        #   "off"   — withheld from the model entirely (default)
        #   "ask"   — offered, but each call waits for per-call user approval
        #   "allow" — offered and run freely
        "action_tool_policy": "off",
        "system_prompt": "",
        "timeout": 120,
        # How long Ollama keeps the model resident in VRAM between messages.
        # "30m"/"-1" (never unload)/"0" (unload now). Avoids cold-reload latency
        # when you return to an idle chat. Empty → Ollama's 5-minute default.
        "keep_alive": "30m",
        # CPU/GPU offload: -1 = Auto (Ollama auto-fits layers to VRAM).
        # 0–100 = percent of model layers to force onto the GPU; the
        # remainder runs on CPU. See OllamaManager.get_model_layers.
        "gpu_offload": -1,
        # Memory curator (the model that extracts facts after each exchange).
        # Empty → same auto-selected model as chat. A dedicated model (e.g.
        # "mistral:latest") gives better extraction but must share VRAM.
        "memory_model": DEFAULT_MEMORY_MODEL,
        # Curator CPU/GPU offload — same scale as gpu_offload above. Default 0
        # (all CPU/RAM): OS-neutral and never evicts the chat model from a small
        # GPU. Boxes with spare VRAM can set -1 (Auto) or a percent to use the GPU.
        # -1 = Auto (let Ollama fit it). The curator is the resident chat model
        # now, so there is nothing to keep off the GPU; pinning it to CPU (0)
        # only bought coexistence with a second, separate curator model.
        "memory_gpu_offload": -1,
        # Similar-fact merge: when a newly extracted fact's embedding is at least
        # this cosine-similar to an existing fact, UPDATE that fact in place
        # instead of appending a duplicate ("edit with new info"). 0 disables
        # (always append). Calibrated on nomic-embed-text: genuine updates
        # (mileage/title/location changes) score 0.81–0.99, while distinct facts
        # top out ~0.61 — so 0.80 catches updates and never merges unrelated
        # facts. Lower to catch looser rephrases; raise toward 1.0 to be stricter.
        "memory_merge_threshold": 0.80,
        # Seconds of quiet before the curator reads a conversation. This is the
        # "conversation is over" signal; each new message restarts the clock.
        "memory_extract_idle": 120,
    }

    def get_settings(self) -> Dict[str, Any]:
        conn = self._connect()
        cur = conn.cursor()
        cur.execute("SELECT key, value FROM settings")
        rows = cur.fetchall()
        conn.close()
        result = dict(self._SETTINGS_DEFAULTS)
        for row in rows:
            try:
                result[row["key"]] = json.loads(row["value"])
            except Exception:
                result[row["key"]] = row["value"]
        return result

    def update_settings(self, data: Dict[str, Any]):
        conn = self._connect()
        try:
            cur = conn.cursor()
            for key, value in data.items():
                if key in self._SETTINGS_DEFAULTS:
                    cur.execute(
                        "INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)",
                        (key, json.dumps(value))
                    )
            conn.commit()
        finally:
            conn.close()

# -----------------------------
# Store Instance
# -----------------------------
from ..nexus_config import MEMORY_DB

DB_PATH = MEMORY_DB
store = PersistentMemoryStore(DB_PATH)