"""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])