core/store.py

"""Aurelius core — persistent graph + vector store.

Backend: SQLite (stdlib, zero infra, runs anywhere the demo runs). The
schema and the public interface are deliberately shaped like the Postgres
+ pgvector layout the project will graduate to (nodes/edges tables, a
reverse index on edges.dst, top-k similarity queries) so the swap is a
config change plus one class, not a redesign. At demo scale (≤ a few
hundred thousand vectors) brute-force numpy similarity beats maintaining
an ANN index anyway; past that, the same method signature is answered by
pgvector's `<->` operator instead.

Embeddings are stored as float32 BLOBs. Two vector columns per node:
  text_emb   — sentence-transformers over the adapter's enriched text
  struct_emb — node2vec over the stored edge list (representation.py)

StoreBackedSource at the bottom is the shared GraphSource implementation
for every ingested-mode adapter (biomed, news, finance): subclass,
set name/description/edge_types, done.
"""

from __future__ import annotations

import json
import sqlite3
import threading
from pathlib import Path
from typing import Iterable, Optional

import numpy as np

from config import AURELIUS_DB
from .source import GraphSource
from .types import Edge, NodeInfo, NodeRef

_SCHEMA = """
CREATE TABLE IF NOT EXISTS nodes (
    source     TEXT NOT NULL,
    id         TEXT NOT NULL,
    title      TEXT NOT NULL,
    text       TEXT DEFAULT '',
    summary    TEXT DEFAULT '',
    features   TEXT DEFAULT '{}',
    text_emb   BLOB,
    struct_emb BLOB,
    PRIMARY KEY (source, id)
);
CREATE INDEX IF NOT EXISTS nodes_title ON nodes(source, title COLLATE NOCASE);

CREATE TABLE IF NOT EXISTS edges (
    source TEXT NOT NULL,
    src    TEXT NOT NULL,
    dst    TEXT NOT NULL,
    type   TEXT NOT NULL DEFAULT 'link',
    weight REAL NOT NULL DEFAULT 1.0,
    PRIMARY KEY (source, src, dst, type)
);
CREATE INDEX IF NOT EXISTS edges_src ON edges(source, src);
CREATE INDEX IF NOT EXISTS edges_dst ON edges(source, dst);
"""


def _to_blob(v: np.ndarray | None) -> bytes | None:
    if v is None:
        return None
    return np.asarray(v, dtype=np.float32).tobytes()


def _from_blob(b: bytes | None) -> np.ndarray | None:
    if b is None:
        return None
    return np.frombuffer(b, dtype=np.float32)


class GraphStore:
    """Thread-safe (single connection + lock) SQLite graph/vector store."""

    def __init__(self, path: str | Path = AURELIUS_DB):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(self.path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._lock = threading.Lock()
        with self._lock:
            self._conn.executescript(_SCHEMA)
            self._conn.commit()

    # ── writes ───────────────────────────────────────────────────────────
    def upsert_nodes(self, source: str, rows: Iterable[dict]):
        """rows: dicts with id, title and optional text/summary/features."""
        with self._lock:
            self._conn.executemany(
                """INSERT INTO nodes(source, id, title, text, summary, features)
                   VALUES(?,?,?,?,?,?)
                   ON CONFLICT(source, id) DO UPDATE SET
                     title=excluded.title,
                     text=CASE WHEN excluded.text != '' THEN excluded.text ELSE nodes.text END,
                     summary=CASE WHEN excluded.summary != '' THEN excluded.summary ELSE nodes.summary END,
                     features=excluded.features""",
                [(source, r["id"], r["title"], r.get("text", ""),
                  r.get("summary", ""), json.dumps(r.get("features", {})))
                 for r in rows])
            self._conn.commit()

    def upsert_edges(self, source: str, rows: Iterable[tuple]):
        """rows: (src_id, dst_id, type, weight) tuples."""
        with self._lock:
            self._conn.executemany(
                """INSERT OR REPLACE INTO edges(source, src, dst, type, weight)
                   VALUES(?,?,?,?,?)""",
                [(source, s, d, t, w) for (s, d, t, w) in rows])
            self._conn.commit()

    def delete_source(self, source: str):
        """Wipe a source's nodes+edges — re-ingests start clean so retired
        edge types / nodes don't linger from a previous run."""
        with self._lock:
            self._conn.execute("DELETE FROM edges WHERE source=?", (source,))
            self._conn.execute("DELETE FROM nodes WHERE source=?", (source,))
            self._conn.commit()

    def set_embeddings(self, source: str, kind: str,
                       embs: dict[str, np.ndarray]):
        col = {"text": "text_emb", "struct": "struct_emb"}[kind]
        with self._lock:
            self._conn.executemany(
                f"UPDATE nodes SET {col}=? WHERE source=? AND id=?",
                [(_to_blob(v), source, k) for k, v in embs.items()])
            self._conn.commit()

    # ── reads ────────────────────────────────────────────────────────────
    def get_node(self, source: str, id: str) -> Optional[dict]:
        with self._lock:
            row = self._conn.execute(
                "SELECT id, title, text, summary, features FROM nodes "
                "WHERE source=? AND id=?", (source, id)).fetchone()
        if not row:
            return None
        return {"id": row[0], "title": row[1], "text": row[2],
                "summary": row[3], "features": json.loads(row[4] or "{}")}

    def find_nodes(self, source: str, query: str, limit: int = 10) -> list[dict]:
        """Exact id → exact title → substring title match, ranked short-first."""
        q = query.strip()
        with self._lock:
            row = self._conn.execute(
                "SELECT id, title FROM nodes WHERE source=? AND id=?",
                (source, q)).fetchone()
            if row:
                return [{"id": row[0], "title": row[1]}]
            row = self._conn.execute(
                "SELECT id, title FROM nodes WHERE source=? AND title=? "
                "COLLATE NOCASE", (source, q)).fetchone()
            if row:
                return [{"id": row[0], "title": row[1]}]
            rows = self._conn.execute(
                "SELECT id, title FROM nodes WHERE source=? AND title LIKE ? "
                "COLLATE NOCASE ORDER BY LENGTH(title) LIMIT ?",
                (source, f"%{q}%", limit)).fetchall()
        return [{"id": r[0], "title": r[1]} for r in rows]

    def suggest_titles(self, source: str, query: str,
                       limit: int = 8) -> list[dict]:
        """Type-ahead lookup: prefix matches (on title or id) rank ahead of
        mid-string matches, shortest title first. Carries features so the
        UI can label each hit by kind. The ingested-mode counterpart to
        Wikipedia's opensearch."""
        q = query.strip()
        if not q:
            return []
        sub = f"%{q}%"
        prefix = f"{q}%"
        with self._lock:
            rows = self._conn.execute(
                """SELECT id, title, features FROM nodes
                   WHERE source=? AND (title LIKE ? COLLATE NOCASE
                                    OR id    LIKE ? COLLATE NOCASE)
                   ORDER BY
                     CASE WHEN title LIKE ? COLLATE NOCASE THEN 0
                          WHEN id    LIKE ? COLLATE NOCASE THEN 1
                          ELSE 2 END,
                     LENGTH(title)
                   LIMIT ?""",
                (source, sub, sub, prefix, prefix, limit)).fetchall()
        return [{"id": r[0], "title": r[1],
                 "features": json.loads(r[2] or "{}")} for r in rows]

    def neighbors(self, source: str, id: str) -> list[tuple[str, str, float]]:
        """→ [(dst_id, type, weight)]"""
        with self._lock:
            rows = self._conn.execute(
                "SELECT dst, type, weight FROM edges WHERE source=? AND src=?",
                (source, id)).fetchall()
        return rows

    def get_edge(self, source: str, src: str, dst: str
                 ) -> Optional[tuple[str, float]]:
        """Strongest (type, weight) between two nodes, or None — the
        evidence lookup behind edge_display."""
        with self._lock:
            row = self._conn.execute(
                "SELECT type, weight FROM edges WHERE source=? AND src=? AND dst=? "
                "ORDER BY weight DESC LIMIT 1", (source, src, dst)).fetchone()
        return row

    def back_neighbors(self, source: str, id: str,
                       limit: int = 500) -> list[tuple[str, str, float]]:
        """→ [(src_id, type, weight)] — answered by the edges_dst index,
        the ingested-mode equivalent of Wikipedia's linkshere."""
        with self._lock:
            rows = self._conn.execute(
                "SELECT src, type, weight FROM edges WHERE source=? AND dst=? "
                "LIMIT ?", (source, id, limit)).fetchall()
        return rows

    def titles_for(self, source: str, ids: list[str]) -> dict[str, str]:
        if not ids:
            return {}
        out: dict[str, str] = {}
        with self._lock:
            for i in range(0, len(ids), 500):
                chunk = ids[i:i + 500]
                marks = ",".join("?" * len(chunk))
                for r in self._conn.execute(
                        f"SELECT id, title FROM nodes WHERE source=? AND id IN ({marks})",
                        (source, *chunk)).fetchall():
                    out[r[0]] = r[1]
        return out

    def all_edges(self, source: str) -> list[tuple[str, str, float]]:
        with self._lock:
            return self._conn.execute(
                "SELECT src, dst, weight FROM edges WHERE source=?",
                (source,)).fetchall()

    def random_nodes(self, source: str, k: int = 2) -> list[dict]:
        with self._lock:
            rows = self._conn.execute(
                "SELECT id, title FROM nodes WHERE source=? "
                "ORDER BY RANDOM() LIMIT ?", (source, k)).fetchall()
        return [{"id": r[0], "title": r[1]} for r in rows]

    def node_count(self, source: str) -> int:
        with self._lock:
            return self._conn.execute(
                "SELECT COUNT(*) FROM nodes WHERE source=?", (source,)).fetchone()[0]

    def edge_count(self, source: str) -> int:
        with self._lock:
            return self._conn.execute(
                "SELECT COUNT(*) FROM edges WHERE source=?", (source,)).fetchone()[0]

    def missing_text_embeddings(self, source: str) -> list[tuple[str, str]]:
        """→ [(id, text-or-title)] for nodes with no text_emb yet."""
        with self._lock:
            rows = self._conn.execute(
                "SELECT id, CASE WHEN text != '' THEN text ELSE title END "
                "FROM nodes WHERE source=? AND text_emb IS NULL",
                (source,)).fetchall()
        return rows

    def embeddings(self, source: str, kind: str = "text"
                   ) -> tuple[list[str], np.ndarray]:
        """All (ids, matrix) for a source — the brute-force ANN workhorse."""
        col = {"text": "text_emb", "struct": "struct_emb"}[kind]
        with self._lock:
            rows = self._conn.execute(
                f"SELECT id, {col} FROM nodes WHERE source=? AND {col} IS NOT NULL",
                (source,)).fetchall()
        if not rows:
            return [], np.zeros((0, 0), dtype=np.float32)
        ids = [r[0] for r in rows]
        mat = np.vstack([_from_blob(r[1]) for r in rows])
        return ids, mat

    def get_embedding(self, source: str, id: str,
                      kind: str = "text") -> Optional[np.ndarray]:
        col = {"text": "text_emb", "struct": "struct_emb"}[kind]
        with self._lock:
            row = self._conn.execute(
                f"SELECT {col} FROM nodes WHERE source=? AND id=?",
                (source, id)).fetchone()
        return _from_blob(row[0]) if row and row[0] else None

    def topk_similar(self, source: str, query: np.ndarray, k: int = 10,
                     kind: str = "text",
                     exclude: set[str] | None = None) -> list[tuple[str, float]]:
        """Top-k cosine neighbours of `query` among a source's vectors.
        Brute-force numpy — the pgvector `ORDER BY emb <-> $1 LIMIT k`
        equivalent at demo scale."""
        ids, mat = self.embeddings(source, kind)
        if not ids:
            return []
        q = np.asarray(query, dtype=np.float32)
        qn = np.linalg.norm(q)
        if qn == 0:
            return []
        norms = np.linalg.norm(mat, axis=1)
        norms[norms == 0] = 1e-9
        sims = (mat @ q) / (norms * qn)
        order = np.argsort(-sims)
        out: list[tuple[str, float]] = []
        excl = exclude or set()
        for i in order:
            if ids[i] in excl:
                continue
            out.append((ids[i], float(sims[i])))
            if len(out) >= k:
                break
        return out


# Shared default store instance (lazy).
_STORE: Optional[GraphStore] = None


def get_store() -> GraphStore:
    global _STORE
    if _STORE is None:
        _STORE = GraphStore()
    return _STORE


# ══════════════════════════════════════════════════════════════
# StoreBackedSource — GraphSource over ingested data
# ══════════════════════════════════════════════════════════════

class StoreBackedSource(GraphSource):
    """Base adapter for ingested-mode sources: subclass, set name/
    description/edge_types, run the matching ingest script, done."""

    supports_backlinks = True   # edges_dst index makes inbound cheap

    def __init__(self):
        self._store: Optional[GraphStore] = None

    @property
    def store(self) -> GraphStore:
        if self._store is None:
            self._store = get_store()
        return self._store

    def ingested(self) -> bool:
        return self.store.node_count(self.name) > 0

    def _ref(self, id: str, title: str | None = None) -> NodeRef:
        if title is None:
            node = self.store.get_node(self.name, id)
            title = node["title"] if node else id
        return NodeRef(source=self.name, id=id, title=title)

    async def resolve(self, query: str) -> Optional[NodeRef]:
        if not self.ingested():
            return None
        hits = self.store.find_nodes(self.name, query)
        if not hits:
            return None
        return self._ref(hits[0]["id"], hits[0]["title"])

    async def neighbors(self, n: NodeRef, *,
                        hunt_id: str | None = None,
                        priority_ids: set[str] | None = None) -> list[Edge]:
        rows = self.store.neighbors(self.name, n.id)
        titles = self.store.titles_for(self.name, [r[0] for r in rows])
        return [Edge(src=n,
                     dst=NodeRef(self.name, dst, titles.get(dst, dst)),
                     type=t, weight=w)
                for (dst, t, w) in rows]

    async def back_neighbors(self, n: NodeRef, limit: int = 500) -> list[Edge]:
        rows = self.store.back_neighbors(self.name, n.id, limit)
        titles = self.store.titles_for(self.name, [r[0] for r in rows])
        return [Edge(src=NodeRef(self.name, src, titles.get(src, src)),
                     dst=n, type=t, weight=w)
                for (src, t, w) in rows]

    async def node_info(self, n: NodeRef, rich: bool = False) -> NodeInfo:
        node = self.store.get_node(self.name, n.id)
        if not node:
            return NodeInfo(text=n.title)
        return NodeInfo(text=node["text"] or node["title"],
                        summary=node["summary"],
                        features=node["features"])

    async def node_infos(self, ns: list[NodeRef]) -> list[NodeInfo]:
        return [await self.node_info(n) for n in ns]

    # ── recommendations ──────────────────────────────────────────────────
    async def suggest(self, query: str, limit: int = 8) -> list[dict]:
        if not self.ingested():
            return []
        out: list[dict] = []
        for r in self.store.suggest_titles(self.name, query, limit):
            feats = r.get("features") or {}
            out.append({
                "id": r["id"], "title": r["title"],
                "kind": feats.get("kind"),
                "subtitle": self.suggest_subtitle(r["id"], r["title"], feats),
            })
        return out

    def suggest_subtitle(self, node_id: str, title: str,
                         features: dict) -> Optional[str]:
        """Short hint shown under a suggestion. Default: the node kind
        (and the id when it differs from the title, e.g. a ticker).
        Adapters override for richer hints."""
        kind = features.get("kind")
        label = kind.replace("_", " ") if kind else None
        show_id = (node_id and node_id != title and len(node_id) <= 8
                   and node_id.lower() not in title.lower())
        if show_id:
            return f"{node_id} · {label}" if label else node_id
        return label

    # ── edge evidence ────────────────────────────────────────────────────
    def format_edge(self, typ: str, weight: float) -> str:
        """Human phrase for a typed edge. Adapters override for domain
        vocabulary; the default just de-snake-cases the type."""
        return typ.replace("_", " ")

    async def edge_display(self, src: NodeRef, dst: NodeRef) -> Optional[str]:
        row = self.store.get_edge(self.name, src.id, dst.id)
        if row is None:
            return None
        return self.format_edge(row[0], row[1])