feat: wire --pool / topk_retrieve override through brain_ask and run_ablation (#92)

CHANGELOG.md

@@ -6,6 +6,8 @@ recorded after the fact rather than backdated to look contemporaneous.
 
 ## Unreleased
 
+**`sgm-ask` and `eval/run_ablation.py`: `--pool` override for the dense candidate pool.** `--pool N` (or `BRAIN_TOPK_RETRIEVE=N`) overrides `TOPK_RETRIEVE` without changing the 60-note default, and `DEFAULT_GRID` adds `topk-15` (`GHOPS=15`) so pool-size sweeps can be compared across vaults. Thanks @Jason-Vaughan (#97), who picked up #92.
+
 ## v0.3.7 - 2026-11-05
 
 **Recall: the cross-encoder rerank is off by default.** `sgm-ask`, the MCP tool and `eval/run_eval.py` now order candidates by dense e5 similarity; `--rerank` or `SGM_RERANK=1` brings the cross-encoder back. Measured before the change: on our 184-question gold set the reranker cut body nDCG@12 from 0.815 to 0.640, and on DolphinBench data (Riley, 50 tests, 1000 sessions) Recall@5 fell from 0.61 to 0.48; it still helped short title-style queries. Dense recall is also faster on CPU (about 0.04 s against 2.8 s per query on that slice). `rerank_candidates()` takes `sims=` and keeps dense order when no cross-encoder is passed.

eval/run_ablation.py

@@ -40,52 +40,61 @@
 from run_eval import SMALL_N, mrr_at, ndcg_at, recall_at  # noqa: E402
 
 
-def _default_cell():
-    return {'hops': 1, 'ghops': ba.GHOPS, 'gmax': ba.GMAX, 'topk': ba.TOPK_RETRIEVE,
+def _default_cell(topk_retrieve=None):
+    return {'hops': 1, 'ghops': ba.GHOPS, 'gmax': ba.GMAX,
+            'topk': ba.TOPK_RETRIEVE if topk_retrieve is None else topk_retrieve,
             'entity_gate': True}
 
 
-def _fill_cell(partial, index):
-    cell = _default_cell()
+def _fill_cell(partial, index, topk_retrieve=None):
+    cell = _default_cell(topk_retrieve=topk_retrieve)
     cell.update(partial)
     if not cell.get('name'):
         cell['name'] = 'cell-%d' % index
     if cell['hops'] not in (0, 1):
         raise ValueError("cell %r: hops must be 0 or 1 (score_grid only runs the 1-hop "
                          "expand_1hop), got %r" % (cell['name'], cell['hops']))
+    if cell['topk'] < 1:
+        raise ValueError("cell %r: topk must be >= 1, got %r" % (cell['name'], cell['topk']))
     return cell
 
 
-DEFAULT_GRID = [
-    _fill_cell({'name': '0-hop', 'hops': 0}, 0),
-    _fill_cell({'name': '1-hop'}, 1),
-    _fill_cell({'name': 'gmax-10', 'gmax': 10}, 2),
-    _fill_cell({'name': 'gmax-20', 'gmax': 20}, 3),
-    _fill_cell({'name': 'gmax-80', 'gmax': 80}, 4),
-    _fill_cell({'name': 'topk-20', 'topk': 20}, 5),
-    _fill_cell({'name': 'topk-80', 'topk': 80}, 6),
-    _fill_cell({'name': 'ghops-5', 'ghops': 5}, 7),
-    _fill_cell({'name': 'ghops-30', 'ghops': 30}, 8),
-    _fill_cell({'name': 'entity-gate-off', 'entity_gate': False}, 9),
+_DEFAULT_CELL_SPECS = [
+    {'name': '0-hop', 'hops': 0},
+    {'name': '1-hop'},
+    {'name': 'gmax-10', 'gmax': 10},
+    {'name': 'gmax-20', 'gmax': 20},
+    {'name': 'gmax-80', 'gmax': 80},
+    {'name': 'topk-15', 'topk': ba.GHOPS},
+    {'name': 'topk-20', 'topk': 20},
+    {'name': 'topk-80', 'topk': 80},
+    {'name': 'ghops-5', 'ghops': 5},
+    {'name': 'ghops-30', 'ghops': 30},
+    {'name': 'entity-gate-off', 'entity_gate': False},
 ]
+DEFAULT_GRID = [_fill_cell(c, i) for i, c in enumerate(_DEFAULT_CELL_SPECS)]
 GRID_PRESETS = {'default': DEFAULT_GRID}
 
 _BOOL_KEYS = {'entity_gate'}
 _STR_KEYS = {'name'}
 
 
-def parse_grid(spec):
+def parse_grid(spec, topk_retrieve=None):
     """'default' | a path or literal JSON array of cell dicts | ';'-separated key=val cells."""
     spec = (spec or 'default').strip()
     if spec in GRID_PRESETS:
-        return GRID_PRESETS[spec]
+        if topk_retrieve is None:
+            return GRID_PRESETS[spec]
+        return [_fill_cell(c, i, topk_retrieve=topk_retrieve)
+                for i, c in enumerate(_DEFAULT_CELL_SPECS)]
 
     path = Path(spec)
     if path.exists():
         spec = path.read_text(encoding='utf-8')
 
     if spec.lstrip().startswith('['):
-        return [_fill_cell(c, i) for i, c in enumerate(json.loads(spec))]
+        return [_fill_cell(c, i, topk_retrieve=topk_retrieve)
+                for i, c in enumerate(json.loads(spec))]
 
     cells = []
     for i, part in enumerate(spec.split(';')):
@@ -101,7 +110,7 @@ def parse_grid(spec):
                 partial[k] = v
             else:
                 partial[k] = int(v)
-        cells.append(_fill_cell(partial, i))
+        cells.append(_fill_cell(partial, i, topk_retrieve=topk_retrieve))
     return cells
 
 
@@ -217,6 +226,8 @@ def main():
     ap.add_argument('--grid', default='default',
                     help="'default', a JSON array of cells, a path to one, or "
                          "';'-separated key=val cells (default: %(default)s)")
+    ap.add_argument('--pool', '--topk-retrieve', dest='topk_retrieve', type=int, default=None,
+                    help='override default dense candidate pool size (topk) for cells that do not set topk')
     ap.add_argument('--list-grid', action='store_true',
                     help='print the named grid cells and exit (no model/index load)')
     ap.add_argument('--out', help='output prefix for <prefix>.csv/.md (default: eval/ablation-<ts>)')
@@ -228,7 +239,13 @@ def main():
     if a.rerank:
         os.environ['SGM_RERANK'] = '1'
 
-    grid = parse_grid(a.grid)
+    pool = a.topk_retrieve
+    if pool is None:
+        pool, _ = ba._parse_pool([])
+    elif pool < 1:
+        return ap.error('--pool / --topk-retrieve must be >= 1')
+
+    grid = parse_grid(a.grid, topk_retrieve=pool)
     if a.list_grid:
         print_grid(grid)
         return 0

src/sqlite_graph_memory/brain_ask.py

@@ -9,11 +9,13 @@
   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
+  python brain_ask.py --pool 15 --graph "q"   # override dense candidate pool size (default: 60)
   python brain_ask.py --ppr --graph "question"  # + PageRank lane in the ORDERING (see ppr.py)
   python brain_ask.py --rerank "question"     # opt in to the cross-encoder rerank (off by default)
 
 Config (env, all optional):
   BRAIN_INDEX_DIR   dir holding the embedding index (default: ./index)
+  BRAIN_TOPK_RETRIEVE  dense candidate pool size before rerank (default: 60; CLI --pool overrides)
   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
@@ -449,12 +451,56 @@ def maybe_star_prompt(stream=None, env=None, interactive=None):
         return False
 
 
+def _parse_pool(args, env=None):
+    """Extract optional `--pool N` / `--topk-retrieve N` (or BRAIN_TOPK_RETRIEVE) -> (int|None, remaining_args)."""
+    env = os.environ if env is None else env
+    pool = None
+    env_val = (env.get('BRAIN_TOPK_RETRIEVE') or '').strip()
+    if env_val:
+        try:
+            pool = int(env_val)
+        except ValueError as exc:
+            sys.stderr.write('invalid BRAIN_TOPK_RETRIEVE=%r (expected positive integer)\n' % env_val)
+            raise SystemExit(2) from exc
+        if pool < 1:
+            sys.stderr.write('invalid BRAIN_TOPK_RETRIEVE=%r (must be >= 1)\n' % env_val)
+            raise SystemExit(2)
+    rest = []
+    i = 0
+    while i < len(args):
+        a = args[i]
+        if a in ('--pool', '--topk-retrieve'):
+            if i + 1 >= len(args):
+                sys.stderr.write('%s requires a positive integer argument\n' % a)
+                raise SystemExit(2)
+            raw = args[i + 1]
+            i += 2
+        elif a.startswith('--pool=') or a.startswith('--topk-retrieve='):
+            flag, _, raw = a.partition('=')
+            a = flag
+            i += 1
+        else:
+            rest.append(a)
+            i += 1
+            continue
+        try:
+            pool = int(raw)
+        except ValueError as exc:
+            sys.stderr.write('invalid %s value %r (expected positive integer)\n' % (a, raw))
+            raise SystemExit(2) from exc
+        if pool < 1:
+            sys.stderr.write('invalid %s value %r (must be >= 1)\n' % (a, raw))
+            raise SystemExit(2)
+    return pool, rest
+
+
 def main():
     args = sys.argv[1:]
     if '--version' in args:
         from sqlite_graph_memory import __version__
         print('sgm-ask %s' % __version__)
         return
+    pool, args = _parse_pool(args)
     ask = '--ask' in args
     graph = '--graph' in args; ab = '--ab' in args
     flags = {'--ask', '--graph', '--ab', '--ppr', '--rerank'}
@@ -471,7 +517,7 @@ def main():
     sims = query_sims(enc, emb, query)
 
     # dense retrieve -> dedup by path (best chunk per file) -> top-K unique files
-    order = retrieve_candidates(sims, meta)
+    order = retrieve_candidates(sims, meta) if pool is None else retrieve_candidates(sims, meta, topk=pool)
 
     # 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

tests/test_brain_ask.py

@@ -217,3 +217,31 @@ def test_missing_embedding_index_reports_error_and_exits_nonzero(tmp_path):
     assert result.stdout == ""
     assert "need an embedding index" in result.stderr
     assert "index_notes.py first" in result.stderr
+
+
+def test_parse_pool_flag_and_env_override_without_changing_default():
+    from sqlite_graph_memory import brain_ask as ba
+
+    assert ba.TOPK_RETRIEVE == 60
+    pool, rest = ba._parse_pool(["--graph", "how does memory work"], env={})
+    assert pool is None
+    assert rest == ["--graph", "how does memory work"]
+
+    pool, rest = ba._parse_pool(["--pool", "15", "--graph", "how does memory work"], env={})
+    assert pool == 15
+    assert rest == ["--graph", "how does memory work"]
+
+    pool, rest = ba._parse_pool(["--topk-retrieve=20", "query"], env={"BRAIN_TOPK_RETRIEVE": "15"})
+    assert pool == 20
+    assert rest == ["query"]
+
+    pool, rest = ba._parse_pool(["query"], env={"BRAIN_TOPK_RETRIEVE": "15"})
+    assert pool == 15
+    assert rest == ["query"]
+
+    with pytest.raises(SystemExit):
+        ba._parse_pool(["--pool", "0"], env={})
+    with pytest.raises(SystemExit):
+        ba._parse_pool(["--pool", "abc"], env={})
+    with pytest.raises(SystemExit):
+        ba._parse_pool(["--pool"], env={})

tests/test_eval.py

@@ -432,3 +432,31 @@ def test_parse_grid_rejects_unsupported_hop_counts():
         parse_grid('hops=2')
     with pytest.raises(ValueError):
         parse_grid('[{"name": "two-hop", "hops": 2}]')
+
+
+def test_pool_override_wires_topk_retrieve_without_changing_default():
+    """Issue #92: --pool / topk_retrieve overrides default cell topk while leaving
+    explicit topk cells and the library default (60) untouched."""
+    default_grid = parse_grid('default')
+    by_name = {c['name']: c for c in default_grid}
+    assert by_name['0-hop']['topk'] == 60
+    assert by_name['1-hop']['topk'] == 60
+    assert by_name['topk-15']['topk'] == 15
+    assert by_name['topk-20']['topk'] == 20
+
+    overridden = parse_grid('default', topk_retrieve=15)
+    by_name_15 = {c['name']: c for c in overridden}
+    assert by_name_15['0-hop']['topk'] == 15
+    assert by_name_15['1-hop']['topk'] == 15
+    assert by_name_15['topk-20']['topk'] == 20
+    assert by_name_15['topk-80']['topk'] == 80
+
+    custom = parse_grid('hops=0,name=vec;hops=1,topk=30,name=explicit', topk_retrieve=15)
+    assert custom[0]['topk'] == 15
+    assert custom[1]['topk'] == 30
+
+    r = subprocess.run([sys.executable, str(ROOT / 'eval' / 'run_ablation.py'),
+                        '--list-grid', '--pool', '15'],
+                       capture_output=True, text=True)
+    assert r.returncode == 0, r.stderr
+    assert '1-hop                 1    15   40   15  True' in r.stdout