sqlite-graph-memory 0.3.0__py3-none-any.whl

sqlite_graph_memory/__init__.py

@@ -0,0 +1,26 @@
+"""Graph RAG on SQLite for AI agents — vectors for entry points, hand-written
+``[[wikilinks]]`` as the graph, cross-encoder rerank on top.
+
+The package ships four command-line entry points rather than a library API,
+because that is how the pilot is actually used:
+
+``sgm-index``
+    Walk a folder of markdown notes, extract ``[[wikilinks]]``, embed the
+    chunks and write everything into one SQLite file.
+``sgm-ask``
+    Answer a question over that file: vector search for entry points, one hop
+    along the wikilink graph, then rerank.
+``sgm-turnstate-hook`` / ``sgm-turnstate-show``
+    Record and inspect per-session agent memory ("what happened last turn").
+
+Importing the submodules directly is supported and stable — the tests do it —
+but the heavy ML dependencies are loaded lazily inside the functions that need
+them, so importing this package costs nothing beyond ``numpy``. Install the
+``embeddings`` extra to get ``sentence-transformers`` and ``torch``; without it
+the indexer and the reranker will tell you what is missing instead of failing
+on an ``ImportError`` at startup.
+"""
+
+__version__ = "0.3.0"
+
+__all__ = ["__version__"]

sqlite_graph_memory/brain_ask.py

@@ -0,0 +1,310 @@
+# -*- coding: utf-8 -*-
+"""brain_ask.py — vector retrieval + graph expansion + A/B harness over a markdown vault.
+
+Pipeline: e5-base dense retrieve (top-60, dedup by file) -> optional 1-hop wikilink
+graph expansion -> cross-encoder rerank -> top-12.
+
+USAGE:
+  python brain_ask.py "question"              # vector-only (Direct memory)
+  python brain_ask.py --graph "question"      # vector + graph expansion (Associative)
+  python brain_ask.py --ab "question"         # run BOTH, diff, log to SQLite (ab_recall)
+  python brain_ask.py --ask "question"        # emit a context bundle for an LLM agent
+
+Config (env, all optional):
+  BRAIN_INDEX_DIR   dir holding the embedding index (default: ./index)
+  TURNSTATE_DB      SQLite db for A/B telemetry (default: ./turnstate.db)
+  BRAIN_ANSWER_OUT  file to mirror the answer into (default: ./_brain_answer.txt)
+  BRAIN_AB_JSON     if set, --ab also writes a structured JSON diff to this path
+"""
+import os, re, sys, pickle, datetime, sqlite3, time
+from pathlib import Path
+import numpy as np
+
+# Make stdout/stderr UTF-8 so non-ASCII notes never crash a Windows console.
+for _s in (sys.stdout, sys.stderr):
+    try: _s.reconfigure(encoding='utf-8')
+    except Exception: pass
+
+INDEX_DIR = Path(os.getenv('BRAIN_INDEX_DIR', 'index'))
+AB_DB = Path(os.getenv('TURNSTATE_DB', 'turnstate.db'))
+EMB = INDEX_DIR / '_brain_e5.npy'
+META = INDEX_DIR / '_brain_e5_meta.pkl'
+E5_MODEL = 'intfloat/multilingual-e5-base'
+RERANK_MODEL = 'cross-encoder/mmarco-mMiniLMv2-L12-H384-v1'
+TOPK_RETRIEVE = 60
+TOPN = 12
+GHOPS, GMAX = 15, 40   # graph-expansion: expand from top-GHOPS hits, add at most GMAX neighbours
+
+# Known single-token product/tool names for the entity gate (see looks_like_entity).
+# Extend with the tools/projects that appear in YOUR vault.
+ENTITY_TOOLS = {'obsidian', 'syncthing', 'sqlite', 'telegram', 'whatsapp', 'n8n'}
+
+
+
+_EMBEDDINGS_HINT = (
+    "This step needs the optional embedding stack (sentence-transformers, torch).\n"
+    "Install it with:  pip install 'sqlite-graph-memory[embeddings]'\n"
+    "The base install stays small on purpose: the graph half, the SQLite schema\n"
+    "and the turnstate memory work without it."
+)
+
+
+def _require_sentence_transformers():
+    """Import sentence_transformers, or exit with instructions.
+
+    Raises SystemExit (not ImportError) so the CLI prints one actionable line
+    rather than a traceback through site-packages.
+    """
+    try:
+        import sentence_transformers  # noqa: F401
+    except ModuleNotFoundError as exc:
+        if exc.name not in ("sentence_transformers", "torch"):
+            raise
+        raise SystemExit(_EMBEDDINGS_HINT) from exc
+    return sentence_transformers
+
+def pick_device():
+    try:
+        import torch
+        return 'cuda' if torch.cuda.is_available() else 'cpu'
+    except Exception:
+        return 'cpu'
+
+
+def looks_like_entity(q):
+    """True if the query is a NAME / tool / project lookup.
+
+    Why: in our A/B runs graph expansion helped THEME queries but hurt ENTITY queries
+    (a name pulls in every note that links to the person's card — mostly noise).
+    CONSERVATIVE: only fires when confident; on doubt returns False so graph stays ON.
+    Themes (lowercase common-noun phrases) never trip it; names/tools (capitalised /
+    CamelCase / snake_case / known tool token) do. Detection is query-side only."""
+    q = (q or '').strip()
+    toks = q.split()
+    if not toks or len(toks) > 5:
+        return False
+    for idx, t in enumerate(toks):
+        core = t.strip('«».,:;"\'()?!')
+        if not core:
+            continue
+        low = core.lower()
+        if low in ENTITY_TOOLS:        return True   # known product/project name
+        if '_' in low:                 return True   # snake_case tool name
+        if len(core) >= 4 and any(c.isupper() for c in core[1:]) and any(c.islower() for c in core):
+            return True                              # CamelCase (NotebookLM, PyTorch)
+        if core[:1].isupper() and (idx > 0 or len(toks) <= 2):
+            if not (core.isupper() and len(core) <= 3):   # ignore bare topic-acronyms (AI, RAG, DAO)
+                return True                          # proper noun (Feynman, Turing)
+    return False
+
+
+FENCED_CODE_RX = re.compile(r'(?:```[\s\S]*?```|~~~[\s\S]*?~~~)')
+WIKILINK_RX = re.compile(r'\[\[([^\]\|#]+)')
+
+def _links_in(path):
+    """Outgoing wikilink TARGETS in a note (alias/heading stripped, code fences ignored)."""
+    try: t = Path(path).read_text(encoding='utf-8', errors='ignore')
+    except Exception: return []
+    t = FENCED_CODE_RX.sub('', t)
+    return [m.strip() for m in WIKILINK_RX.findall(t)]
+
+
+
+# --------------------------------------------------------------------------- pipeline
+# The retrieval pipeline lives in these functions so that anything measuring it
+# (eval/run_eval.py) runs THE SAME CODE the agent runs, instead of a re-implementation
+# that quietly drifts from it. main() below is a CLI over them.
+
+def load_index(emb_path=None, meta_path=None):
+    """-> (embedding matrix, meta list). Defaults to the index BRAIN_INDEX_DIR points at."""
+    meta = pickle.loads(Path(meta_path or META).read_bytes())
+    emb = np.load(emb_path or EMB)
+    return emb, meta
+
+
+def load_encoder(dev=None):
+    _require_sentence_transformers()
+    from sentence_transformers import SentenceTransformer
+    return SentenceTransformer(E5_MODEL, device=dev or pick_device())
+
+
+def load_reranker(dev=None):
+    _require_sentence_transformers()
+    from sentence_transformers import CrossEncoder
+    return CrossEncoder(RERANK_MODEL, device=dev or pick_device())
+
+
+def query_sims(enc, emb, query):
+    """Cosine similarity of one query against every chunk. e5 wants the 'query: ' prefix."""
+    qv = enc.encode(['query: ' + query], normalize_embeddings=True,
+                    convert_to_numpy=True)[0].astype('float32')
+    return emb @ qv
+
+
+def retrieve_candidates(sims, meta, topk=TOPK_RETRIEVE):
+    """Dense retrieve -> best chunk per file -> top-K unique files (as chunk indices)."""
+    order, seen = [], set()
+    for i in np.argsort(-sims):
+        m = meta[i]
+        if m['path'] in seen: continue
+        seen.add(m['path']); order.append(int(i))
+        if len(order) >= topk: break
+    return order
+
+
+def index_by_basename(meta):
+    """basename(no .md) -> meta indices, used to resolve [[wikilink]] targets to chunks."""
+    by_base = {}
+    for j, mm in enumerate(meta):
+        by_base.setdefault(Path(mm['path']).stem.lower(), []).append(j)
+    return by_base
+
+
+def expand_1hop(sims, meta, base_order, by_base=None, ghops=GHOPS, gmax=GMAX):
+    """Follow 1-hop OUTGOING wikilinks from the top vector hits and return the neighbour
+    chunks to ADD to the candidate set; the reranker decides whether they are relevant.
+
+    Bounded by ghops/gmax. CRASH-SAFETY: any failure returns [] and the caller degrades to
+    vector-only -- worst case you get the old result, never an empty or broken recall."""
+    added = []
+    if not base_order: return added
+    try:
+        by_base = by_base or index_by_basename(meta)
+        in_order = set(base_order)
+        for i in base_order[:ghops]:
+            for tgt in _links_in(meta[i]['path']):
+                idxs = by_base.get(tgt.lower())
+                if not idxs: continue
+                best = max(idxs, key=lambda j: sims[j])   # best chunk of neighbour for THIS query
+                if best in in_order: continue
+                in_order.add(best); added.append(best)
+                if len(added) >= gmax: break
+            if len(added) >= gmax: break
+    except Exception as e:
+        sys.stderr.write(("graph-expansion fell back to vector (%s)" % e) + chr(10))
+        return []
+    return added
+
+
+def rerank_candidates(ce, query, cand, meta, topn=TOPN):
+    """Cross-encoder rerank of candidate chunks -> [(chunk index, score)], best first."""
+    pairs = [(query, meta[i]['title'] + '. ' + meta[i]['snippet']) for i in cand]
+    sc = ce.predict(pairs) if len(pairs) else []
+    return sorted(zip(cand, sc), key=lambda x: -x[1])[:topn]
+
+
+def main():
+    args = sys.argv[1:]
+    ask = '--ask' in args
+    graph = '--graph' in args; ab = '--ab' in args
+    flags = {'--ask', '--graph', '--ab'}
+    query = ' '.join(a for a in args if a not in flags)
+    if not query or not EMB.exists():
+        print('need an embedding index (%s) and a query — run index_notes.py first' % EMB); return
+
+    dev = pick_device()
+    emb, meta = load_index()
+    enc = load_encoder(dev)
+    sims = query_sims(enc, emb, query)
+
+    # dense retrieve -> dedup by path (best chunk per file) -> top-K unique files
+    order = retrieve_candidates(sims, meta)
+
+    # Graph expansion: follow 1-hop OUTGOING wikilinks from the top vector hits, pull in
+    # those neighbour notes (resolved to already-indexed chunks), then let the reranker
+    # decide if they're relevant. Bounded + opt-in (--graph) so base recall is untouched.
+    # Serendipity that the reranker keeps honest.
+    # CRASH-SAFETY: wrapped so ANY failure degrades to vector-only — worst case you get
+    # the old result, never an empty/broken recall.
+    # Graph expansion: 1-hop outgoing wikilinks from the top vector hits, bounded and
+    # opt-in (--graph / --ab) so base recall is untouched. Serendipity the reranker keeps honest.
+    base_order = list(order); added = []
+    if (graph or ab) and order:
+        added = expand_1hop(sims, meta, base_order)
+    graph_order = base_order + added; graph_set = set(added)
+
+    # rerank on the matched chunk (precise)
+    ce = load_reranker(dev)
+    def rerank(cand):
+        return rerank_candidates(ce, query, cand, meta)
+
+    # --ab: run BOTH paths in ONE process, diff, log. Answers "does graph REALLY help?"
+    if ab:
+        A = rerank(base_order); B = rerank(graph_order)
+        a_paths = {meta[i]['path'] for i, _ in A}
+        promoted = [(r, meta[i]['title']) for r, (i, _) in enumerate(B, 1)
+                    if i in graph_set]                        # graph notes that reached top-N
+        new_in_B = [(r, meta[i]['title']) for r, (i, _) in enumerate(B, 1)
+                    if meta[i]['path'] not in a_paths]        # any note B surfaced that A missed
+        out = [f'A/B RECALL: {query}',
+               f'(Direct={len(base_order)} vector / Associative=+{len(added)} graph; '
+               f'Associative promoted {len(new_in_B)} new notes into top-{TOPN}, {len(promoted)} of them via graph)', '',
+               '--- DIRECT memory (vector) ---']
+        for r, (i, sc) in enumerate(A, 1):
+            out.append('%2d. [%6.2f] %s' % (r, float(sc), meta[i]['title']))
+        out += ['', '--- ASSOCIATIVE memory (vector+graph) ---']
+        for r, (i, sc) in enumerate(B, 1):
+            mk = ('  <== +graph' if i in graph_set else
+                  ('  <== NEW' if meta[i]['path'] not in a_paths else ''))
+            out.append('%2d. [%6.2f] %s%s' % (r, float(sc), meta[i]['title'], mk))
+        text = '\n'.join(out)
+        try:
+            con = sqlite3.connect(str(AB_DB), timeout=3.0)
+            con.execute("""CREATE TABLE IF NOT EXISTS ab_recall(
+                id INTEGER PRIMARY KEY AUTOINCREMENT, ts TEXT, query TEXT,
+                cand_added INTEGER, new_in_top INTEGER, promoted_via_graph INTEGER,
+                promoted_titles TEXT, source TEXT)""")
+            con.execute("INSERT INTO ab_recall(ts,query,cand_added,new_in_top,promoted_via_graph,promoted_titles,source)"
+                        " VALUES(?,?,?,?,?,?,?)",
+                        (time.strftime('%Y-%m-%d %H:%M:%S'), query, len(added), len(new_in_B),
+                         len(promoted), '; '.join(t for _, t in promoted)[:1000],
+                         os.getenv('AB_SOURCE', 'work')))
+            con.commit(); con.close()
+        except Exception:
+            pass
+        jpath = os.getenv('BRAIN_AB_JSON')   # structured A/B for a dashboard
+        if jpath:
+            import json as _json
+            def _ser(lst):
+                return [{'rank': r, 'title': meta[i]['title'], 'score': round(float(sc), 2),
+                         'snippet': meta[i].get('snippet', '')[:200],
+                         'graph': i in graph_set,
+                         'new': meta[i]['path'] not in a_paths}
+                        for r, (i, sc) in enumerate(lst, 1)]
+            try:
+                Path(jpath).write_text(_json.dumps(
+                    {'query': query, 'added': len(added), 'new_in_top': len(new_in_B),
+                     'direct': _ser(A), 'associative': _ser(B)}, ensure_ascii=False), encoding='utf-8')
+            except Exception:
+                pass
+        _out = Path(os.getenv('BRAIN_ANSWER_OUT', '_brain_answer.txt'))
+        _out.write_text(text, encoding='utf-8'); print(text)
+        return
+
+    # GATE: wiki-graph helps THEME queries, hurts ENTITY/name queries -> skip graph for entities.
+    gated = graph and looks_like_entity(query)
+    if gated:
+        sys.stderr.write("wiki-graph gated OFF (entity-like query) -> Direct only\n")
+    order = graph_order if (graph and not gated) else base_order
+    reranked = rerank(order)
+    scope = 'all' + (' /graph+%d' % len(graph_set) if graph else '')
+    lines = [f'QUERY: {query}', f'(scope={scope}; e5 {len(order)} files -> rerank; {len(meta)} chunks)', '']
+    if ask:
+        lines.append('=== CONTEXT BUNDLE (paste to your agent) ===\n')
+        for i, sc in reranked[:8]:
+            m = meta[i]
+            hdr = "## %s  [%s] rr=%.2f%s" % (m['title'], m.get('date', ''), float(sc), ' +graph' if i in graph_set else '')
+            lines += [hdr, m['snippet'], '']
+    else:
+        for rank, (i, sc) in enumerate(reranked, 1):
+            m = meta[i]
+            gmark = ' +graph' if i in graph_set else ''
+            lines.append("%2d. [rr=%6.2f] [%s] %s%s" % (rank, float(sc), m.get('date', ''), m['title'], gmark))
+            lines.append("     %s" % m['snippet'][:140])
+    _out = Path(os.getenv('BRAIN_ANSWER_OUT', '_brain_answer.txt'))
+    _out.write_text('\n'.join(lines), encoding='utf-8')
+    print(f'{len(reranked)} hits (scope={scope}) -> {_out}')
+
+
+if __name__ == '__main__':
+    main()

sqlite_graph_memory/index_notes.py

@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+"""index_notes.py — minimal embedding indexer for a folder of markdown notes.
+
+Builds the dense index that brain_ask.py retrieves from:
+  index/_brain_e5.npy       float32 matrix, one row per chunk (L2-normalized)
+  index/_brain_e5_meta.pkl  list of dicts: {path, title, snippet, date}
+
+Chunking is deliberately simple (fixed-size windows over the note body) — the
+cross-encoder reranker downstream forgives coarse chunk boundaries.
+
+USAGE:
+  python index_notes.py <notes_dir>
+Config (env, optional):
+  BRAIN_INDEX_DIR     output dir (default: ./index)
+  BRAIN_INDEX_HIDDEN  1 = also index dot-directories (default: skip them,
+                      so .stversions/.obsidian/.git and Syncthing conflict
+                      copies never enter the index)
+"""
+import os, re, sys, pickle
+from pathlib import Path
+import numpy as np
+
+E5_MODEL = 'intfloat/multilingual-e5-base'
+CHUNK_CHARS = 1200
+CHUNK_OVERLAP = 200
+SNIPPET_CHARS = 600
+
+FM_RX = re.compile(r'^---\r?\n(.*?)\r?\n---\r?\n', re.S)
+
+
+
+_EMBEDDINGS_HINT = (
+    "This step needs the optional embedding stack (sentence-transformers, torch).\n"
+    "Install it with:  pip install 'sqlite-graph-memory[embeddings]'\n"
+    "The base install stays small on purpose: the graph half, the SQLite schema\n"
+    "and the turnstate memory work without it."
+)
+
+
+def _require_sentence_transformers():
+    """Import sentence_transformers, or exit with instructions.
+
+    Raises SystemExit (not ImportError) so the CLI prints one actionable line
+    rather than a traceback through site-packages.
+    """
+    try:
+        import sentence_transformers  # noqa: F401
+    except ModuleNotFoundError as exc:
+        if exc.name not in ("sentence_transformers", "torch"):
+            raise
+        raise SystemExit(_EMBEDDINGS_HINT) from exc
+    return sentence_transformers
+
+def chunks_of(text):
+    step = CHUNK_CHARS - CHUNK_OVERLAP
+    for start in range(0, max(len(text), 1), step):
+        c = text[start:start + CHUNK_CHARS].strip()
+        if len(c) >= 40:
+            yield c
+        if start + CHUNK_CHARS >= len(text):
+            break
+
+
+def fm_date(fm):
+    m = re.search(r'(?m)^date[^:]*:\s*"?(\d{4}-\d{2}-\d{2})', fm)
+    return m.group(1) if m else ''
+
+
+# A synced Obsidian vault is full of markdown that is not a note. `.stversions/`
+# holds Syncthing's old revisions, `*.sync-conflict-*` are its conflict copies,
+# `.obsidian/` is plugin cache and `.git/` is repository internals. Indexing those
+# is worse than noise: a stale revision of a note is semantically almost identical
+# to the live one, so it sits next to it in the top-K and answers with old content.
+CONFLICT_MARKER = 'sync-conflict'
+
+
+def iter_notes(notes_dir):
+    """Files that count as a note under ``notes_dir``.
+
+    Skips dot-directories and Syncthing conflict copies. Set ``BRAIN_INDEX_HIDDEN=1``
+    to index hidden folders too — someone will legitimately keep notes in one.
+    """
+    root = Path(notes_dir)
+    include_hidden = os.getenv('BRAIN_INDEX_HIDDEN', '') not in ('', '0', 'false', 'no')
+    out = []
+    for p in root.rglob('*.md'):
+        rel = p.relative_to(root)
+        if not include_hidden and any(part.startswith('.') for part in rel.parts[:-1]):
+            continue
+        if CONFLICT_MARKER in p.name:
+            continue
+        out.append(p)
+    return sorted(out)
+
+
+def read_note(path):
+    """Text of one note, or None if it cannot be read.
+
+    ``utf-8-sig`` strips the BOM. Reading BOM'd files as plain utf-8 leaves a
+    leading U+FEFF, so the first line is not ``---`` and the frontmatter regex
+    below silently sees a note as having none — losing its date and title.
+    """
+    try:
+        return Path(path).read_text(encoding='utf-8-sig', errors='replace')
+    except Exception:
+        return None
+
+
+def main():
+    if len(sys.argv) < 2:
+        print(__doc__); return
+    notes_dir = Path(sys.argv[1])
+    out_dir = Path(os.getenv('BRAIN_INDEX_DIR', 'index'))
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    meta, texts = [], []
+    for p in iter_notes(notes_dir):
+        t = read_note(p)
+        if t is None:
+            continue
+        m = FM_RX.match(t)
+        fm = m.group(1) if m else ''
+        body = t[m.end():] if m else t
+        date = fm_date(fm)
+        title = p.stem
+        for c in chunks_of(body):
+            meta.append({'path': str(p), 'title': title,
+                         'snippet': c[:SNIPPET_CHARS], 'date': date})
+            texts.append('passage: ' + c)     # e5 expects the passage: prefix
+    if not texts:
+        print('no indexable .md notes found under', notes_dir); return
+
+    _require_sentence_transformers()
+    from sentence_transformers import SentenceTransformer
+    try:
+        import torch
+        dev = 'cuda' if torch.cuda.is_available() else 'cpu'
+    except Exception:
+        dev = 'cpu'
+    enc = SentenceTransformer(E5_MODEL, device=dev)
+    emb = enc.encode(texts, normalize_embeddings=True, convert_to_numpy=True,
+                     batch_size=64, show_progress_bar=True).astype('float32')
+
+    np.save(out_dir / '_brain_e5.npy', emb)
+    (out_dir / '_brain_e5_meta.pkl').write_bytes(pickle.dumps(meta))
+    files = len({m['path'] for m in meta})
+    print('indexed %d chunks from %d files -> %s' % (len(meta), files, out_dir))
+
+
+if __name__ == '__main__':
+    main()