From a8374fd593106343c28b8b151895b47a3a3f7881 Mon Sep 17 00:00:00 2001 From: Raphael Southall <69826455+raphasouthall@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:57:24 +0100 Subject: [PATCH] perf(search): take the write and the 643-entity fan-out off the search path vault_search took 7 to 227 seconds and MCP clients gave up at 30. Entity extraction matched substrings anywhere, so a 4-word query pulled 643 entities and every stage after it was sized by that number. The co-occurrence boost then ran two lookups per entity, 1286 round trips. And every search wrote reinforcement inline, which serialized overlapping searches on the single SQLite writer and grew the table 250k rows a day. Extraction is now one statement, anchored at a word start, capped at 50 and keeping the least frequent entities: a term filling hundreds of triple slots boosts every candidate equally and separates none of them. The boost lookup is one chunked IN query. Reinforcement buffers in memory and drains on a background thread, so the search call issues no writes at all. Co-occurrence stays: its read path feeds the ranking boost and the attractor's clustering blend, and accumulated values are untouched. Three consecutive searches go from 41s, 7.4s, 226.8s to 0.61s, 0.63s, 0.77s. Four threads of overlapping searches no longer serialize. Closes #120 --- src/neurostack/cli/__init__.py | 2 + src/neurostack/cli/search.py | 15 +- src/neurostack/cooccurrence.py | 201 +++++++++++++++++++- src/neurostack/search.py | 175 ++++++++++++++---- tests/conftest.py | 25 +++ tests/test_cooccurrence.py | 191 +++++++++++++++++++ tests/test_search.py | 328 ++++++++++++++++++++++++++++++++- 7 files changed, 892 insertions(+), 45 deletions(-) diff --git a/src/neurostack/cli/__init__.py b/src/neurostack/cli/__init__.py index 878f52f..735cf21 100644 --- a/src/neurostack/cli/__init__.py +++ b/src/neurostack/cli/__init__.py @@ -864,6 +864,8 @@ def main(): p.add_argument("--limit", type=int, default=20, help="Number of top pairs to show (default: 20)") p.add_argument("--json", action="store_true", help="Output as JSON") + p.add_argument("--flush", action="store_true", + help="Write this process's pending reinforcement buffer first") p.set_defaults(func=cmd_cooccurrence) # harvest diff --git a/src/neurostack/cli/search.py b/src/neurostack/cli/search.py index 0e55c78..b989165 100644 --- a/src/neurostack/cli/search.py +++ b/src/neurostack/cli/search.py @@ -1099,10 +1099,23 @@ def cmd_decay(args): def cmd_cooccurrence(args): """Inspect entity co-occurrence pairs.""" - from ..cooccurrence import get_cooccurrence_stats, get_top_pairs + from ..cooccurrence import ( + flush_reinforcement, + get_cooccurrence_stats, + get_top_pairs, + ) from ..schema import DB_PATH, get_db conn = get_db(DB_PATH) + + if args.flush: + # Drains THIS process's pending reinforcement buffer (issue #120). The + # buffer is per-process, so this is the escape hatch for a script or + # REPL that ran searches in-process; the long-lived MCP server drains + # its own buffer on its flush threshold and at shutdown. + written = flush_reinforcement(conn) + print(f"Flushed {written} pending reinforcement pairs.") + pairs = get_top_pairs(conn, limit=args.limit) if args.json: diff --git a/src/neurostack/cooccurrence.py b/src/neurostack/cooccurrence.py index ea17839..e9502cf 100644 --- a/src/neurostack/cooccurrence.py +++ b/src/neurostack/cooccurrence.py @@ -15,8 +15,10 @@ rebuild while either signal is positive. """ +import atexit import logging import sqlite3 +import threading from collections import defaultdict from datetime import datetime, timezone @@ -25,6 +27,12 @@ # Caps the reinforcement (usage) signal; structural weights are raw counts MAX_COOCCURRENCE_WEIGHT = 100.0 +# SQLite caps the number of bound variables per statement (999 by default), so +# every batch query over an arbitrary-length entity/pair list is issued in +# chunks of this size. Shared with the search-side co-occurrence lookup so both +# sides stay under the same limit (issue #120). +SQL_PARAM_CHUNK = 500 + def reinforce_cooccurrence( conn: sqlite3.Connection, entity_pairs: list[tuple[str, str]] @@ -60,9 +68,9 @@ def reinforce_cooccurrence( # Batch-fetch existing reinforcement in a single query canonical_list = sorted(canonical) existing: dict[tuple[str, str], float] = {} - # SQLite has a variable limit; process in chunks of 500 pairs - for chunk_start in range(0, len(canonical_list), 500): - chunk = canonical_list[chunk_start:chunk_start + 500] + # SQLite has a variable limit; process in pair chunks + for chunk_start in range(0, len(canonical_list), SQL_PARAM_CHUNK): + chunk = canonical_list[chunk_start:chunk_start + SQL_PARAM_CHUNK] where_clauses = " OR ".join( "(entity_a = ? AND entity_b = ?)" for _ in chunk ) @@ -100,6 +108,193 @@ def reinforce_cooccurrence( return 0 +# ── Deferred reinforcement (issue #120) ── +# +# Reinforcement is a WRITE and SQLite allows exactly one writer, so doing it +# inside a search put the request on the write lock: two overlapping searches +# serialized, and the second was measured blocking for minutes on a 12.8M-row +# table. Nothing reads the reinforcement column back until the *next* search's +# ranking blend, so the write does not belong on the request path at all. +# Searches buffer pairs in memory and the write happens elsewhere, through +# reinforce_cooccurrence, whose semantics are untouched. +# +# "Elsewhere" is a short-lived background thread with its own connection, not an +# inline flush every N searches: an inline flush would put the write back on one +# request in N, and the request path has to be read-only. Under WAL a concurrent +# writer does not block the readers a search performs. +# +# The buffer is per-process, and that is exactly what makes it work: the MCP +# server is a long-lived process, so pairs buffered by one search are written by +# the drain a later search triggers, by `neurostack cooccurrence --flush`, or by +# the atexit hook when the process stops. A short-lived CLI invocation flushes +# on exit. +# +# Pairs are deduplicated while buffered. A pair repeated across searches inside +# one flush window is therefore reinforced once rather than once per search — +# deliberate: three identical consecutive searches used to triple-bump the same +# pairs, which is query repetition, not association strength. +REINFORCEMENT_FLUSH_THRESHOLD = 5000 + +# Hard ceiling so a persistently failing flush (locked database, disk full) +# cannot grow the buffer without bound. Past the ceiling new pairs are dropped: +# reinforcement is a soft ranking signal, and losing some of it is strictly +# better than an unbounded process. +REINFORCEMENT_BUFFER_MAX = 50_000 + +_reinforcement_lock = threading.Lock() +_reinforcement_buffer: set[tuple[str, str]] = set() +# Path of the database the buffered pairs belong to, so the writer thread and the +# atexit hook can open their own connections instead of using one across threads +# (sqlite3 forbids that). Empty for in-memory databases, which no other thread +# can reach — those pairs wait for an explicit flush_reinforcement call. +_reinforcement_db_path: str = "" +_flush_thread_lock = threading.Lock() +_flush_thread: threading.Thread | None = None + + +def _database_path(conn: sqlite3.Connection) -> str: + """File path backing *conn*'s main database ("" for in-memory).""" + try: + row = conn.execute("PRAGMA database_list").fetchone() + return (row[2] or "") if row else "" + except Exception: + return "" + + +def _flush_in_background(db_path: str) -> threading.Thread | None: + """Drain the buffer on a worker thread with its own connection. + + One writer at a time: while a drain is in flight, later callers just keep + buffering, and whatever they add is picked up by the next drain. + """ + global _flush_thread + with _flush_thread_lock: + if _flush_thread is not None and _flush_thread.is_alive(): + return _flush_thread + thread = threading.Thread( + target=_flush_worker, + args=(db_path,), + name="neurostack-reinforce", + daemon=True, + ) + _flush_thread = thread + thread.start() + return thread + + +def _flush_worker(db_path: str) -> None: + """Open a writer connection, drain the buffer, close it again.""" + try: + conn = sqlite3.connect(db_path, timeout=60.0) + conn.row_factory = sqlite3.Row + try: + conn.execute("PRAGMA busy_timeout=60000") + flush_reinforcement(conn) + finally: + conn.close() + except Exception: + log.debug("background reinforcement flush failed", exc_info=True) + + +def buffer_reinforcement( + conn: sqlite3.Connection, entity_pairs: list[tuple[str, str]] +) -> int: + """Queue entity pairs for a later reinforcement write. + + Writes nothing itself. Once the buffer passes REINFORCEMENT_FLUSH_THRESHOLD a + background thread drains it, so the caller — a search serving a request — + stays read-only. Pairs buffered from an in-memory database (no file another + thread could open) wait for an explicit flush_reinforcement instead. + + Never raises: a caller in the middle of serving a search must never see a + reinforcement failure. Returns the number of pairs still buffered. + """ + if not entity_pairs: + return 0 + + global _reinforcement_db_path + try: + # Canonicalized outside the lock (entity_a < entity_b, the order the + # table stores) so the buffer dedups a pair seen in either direction and + # the critical section stays a single set update. + canonical = {(min(a, b), max(a, b)) for a, b in entity_pairs if a != b} + if not canonical: + return 0 + db_path = "" if _reinforcement_db_path else _database_path(conn) + + with _reinforcement_lock: + _reinforcement_db_path = _reinforcement_db_path or db_path + db_path = _reinforcement_db_path + headroom = REINFORCEMENT_BUFFER_MAX - len(_reinforcement_buffer) + if headroom > 0: + _reinforcement_buffer.update(list(canonical)[:headroom]) + pending = len(_reinforcement_buffer) + + if pending > REINFORCEMENT_FLUSH_THRESHOLD and db_path: + _flush_in_background(db_path) + return pending + except Exception: + log.debug("buffer_reinforcement failed silently", exc_info=True) + return 0 + + +def flush_reinforcement(conn: sqlite3.Connection) -> int: + """Write every buffered pair and clear the buffer. + + The batch is taken out of the buffer before the write, so a concurrent + search keeps buffering into an empty set rather than waiting on the write. + A failed write drops its batch (reinforce_cooccurrence swallows the error + and returns 0) instead of re-queuing it: retrying a batch that failed on a + locked database is how a bounded buffer turns into an unbounded one. + + Returns the number of pairs written. + """ + with _reinforcement_lock: + if not _reinforcement_buffer: + return 0 + batch = list(_reinforcement_buffer) + _reinforcement_buffer.clear() + return reinforce_cooccurrence(conn, batch) + + +def reinforcement_buffer_size() -> int: + """Number of entity pairs waiting to be written.""" + with _reinforcement_lock: + return len(_reinforcement_buffer) + + +def _flush_reinforcement_at_exit() -> None: + """Drain the buffer on interpreter shutdown so the signal is not lost.""" + # A drain already running holds a batch this function cannot see; give it a + # moment to land before writing the rest. The thread is a daemon, so without + # the join its batch would die with the interpreter. + with _flush_thread_lock: + thread = _flush_thread + if thread is not None and thread.is_alive(): + thread.join(timeout=10.0) + + with _reinforcement_lock: + pending = len(_reinforcement_buffer) + db_path = _reinforcement_db_path + if not pending or not db_path: + return + try: + # A fresh connection: the one that buffered these pairs may belong to + # another thread, and sqlite3 forbids cross-thread use. Raw connect + # rather than get_db -- shutdown is no time to run migrations. + conn = sqlite3.connect(db_path, timeout=30.0) + conn.row_factory = sqlite3.Row + try: + flush_reinforcement(conn) + finally: + conn.close() + except Exception: + log.debug("atexit reinforcement flush failed", exc_info=True) + + +atexit.register(_flush_reinforcement_at_exit) + + def persist_cooccurrence(conn: sqlite3.Connection) -> int: """Compute and persist structural co-occurrence weights from triples. diff --git a/src/neurostack/search.py b/src/neurostack/search.py index 2a1da50..a5b01f4 100644 --- a/src/neurostack/search.py +++ b/src/neurostack/search.py @@ -20,7 +20,7 @@ log = logging.getLogger("neurostack") -from .cooccurrence import reinforce_cooccurrence +from .cooccurrence import SQL_PARAM_CHUNK, buffer_reinforcement from .embedder import ( blob_to_embedding, cosine_similarity_batch, @@ -45,6 +45,24 @@ # weak/exploratory query that hit the least-bad target once. PREDICTION_ERROR_MIN_OCCURRENCES = 2 +# Cap on the number of query-matched entities fed to the co-occurrence stage. +# Every stage downstream is sized by this count — the boost lookup joins on it +# and reinforcement pairs it against every result entity — but the cap is not +# there to make a slow query fast. An entity that matches a short query +# alongside hundreds of others carries no ranking signal: a boost that every +# candidate note earns discriminates between none of them. When the cap bites we +# keep the entities appearing in the FEWEST triples, because a term that is +# everywhere in the graph separates nothing while a rare term is precisely the +# one whose associations are worth following (issue #120). +MAX_QUERY_ENTITIES = 50 + +# Cap on result-note entities paired against query entities for reinforcement. +# The pair count is the product of the two caps: 50 x 40 = 2000 pairs per search, +# worst case. Unbounded fan-out is what grew entity_cooccurrence to 12.8M rows +# for 695 notes. Entities are taken in result rank order, so the top-ranked +# note's associations are the ones that survive the cap. +MAX_RESULT_ENTITIES = 40 + def log_prediction_error( conn: sqlite3.Connection, @@ -714,6 +732,95 @@ def decay_hours_since(now=None) -> float | None: ) +def _extract_query_entities(conn: sqlite3.Connection, query: str) -> set[str]: + """Entities in the triples graph that the query's words name. + + One statement for the whole query, not one per word: the per-word version + rescanned the triples table for every term and unioned the results in Python + (issue #120). + + Matching is anchored at a word start — the entity begins with the query word, + or a word inside it does. The old rule was a bare ``LIKE '%word%'``, + substring-anywhere, so "agent" matched most of the graph ("subagent", + "user_agent_string"): the 4-word query "azure foundry knowledge agent" pulled + 643 entities, which are not the query's entities, and every later stage was + sized by that number. Anchoring keeps the variants a user means ("agents", + "agentic", "Azure Foundry") and drops the accidental interior hits. Hyphens + and underscores normalize to spaces on both sides, so "azure-foundry" is + still two words. + + Words of 2 characters or fewer are skipped (they match too much to mean + anything) and the result is capped at MAX_QUERY_ENTITIES, least-frequent + first. + """ + query_words = [w.lower() for w in query.split() if len(w) > 2] + if not query_words: + return set() + + params: list[str] = [] + for word in query_words: + # LIKE wildcards inside a query word would match everything; escape + # them. '_' and '-' are normalized away before escaping. + w = word.replace("-", " ").replace("_", " ") + w = w.replace("\\", r"\\").replace("%", r"\%") + params.extend((f"{w}%", f"% {w}%")) + clauses = " OR ".join( + r"norm LIKE ? ESCAPE '\' OR norm LIKE ? ESCAPE '\'" for _ in query_words + ) + # occurrences = how many triple slots the entity fills, i.e. how common it + # is; ordering by it ascending means the cap keeps the selective entities. + rows = conn.execute( + "SELECT entity, COUNT(*) AS occurrences FROM (" + " SELECT subject AS entity," + " REPLACE(REPLACE(LOWER(subject), '-', ' '), '_', ' ') AS norm" + " FROM triples" + " UNION ALL" + " SELECT object AS entity," + " REPLACE(REPLACE(LOWER(object), '-', ' '), '_', ' ') AS norm" + " FROM triples" + f") WHERE {clauses} " + "GROUP BY entity ORDER BY occurrences ASC, entity ASC LIMIT ?", + [*params, MAX_QUERY_ENTITIES], + ).fetchall() + return {r[0] for r in rows} + + +def _cooccurring_entities( + conn: sqlite3.Connection, query_entities: set[str] +) -> dict[str, float]: + """Entities associated with *query_entities*, mapped to their blended weight. + + ``weight + reinforcement`` (issue #60): structural co-occurrence blended with + the accumulated usage signal, keeping the strongest association per entity. + Entities that are themselves query entities are excluded — an entity the + query already matched is a direct hit, not an association. + + One statement per SQL_PARAM_CHUNK entities rather than two per entity (issue + #120): 643 query entities meant 1286 round trips, 3.5s of the request spent + in fetchall. The result is identical to the per-entity loop — each row is + folded in from whichever side matched, and max() makes the rows a chunk + boundary returns twice idempotent. + """ + cooc_entities: dict[str, float] = {} + qe_list = sorted(query_entities) + for start in range(0, len(qe_list), SQL_PARAM_CHUNK): + chunk = qe_list[start:start + SQL_PARAM_CHUNK] + marks = ",".join("?" * len(chunk)) + rows = conn.execute( + f"SELECT entity_a, entity_b, weight + reinforcement AS w " + f"FROM entity_cooccurrence " + f"WHERE entity_a IN ({marks}) OR entity_b IN ({marks})", + chunk + chunk, + ).fetchall() + for r in rows: + ea, eb, w = r[0], r[1], r[2] + if ea in query_entities and eb not in query_entities: + cooc_entities[eb] = max(cooc_entities.get(eb, 0), w) + if eb in query_entities and ea not in query_entities: + cooc_entities[ea] = max(cooc_entities.get(ea, 0), w) + return cooc_entities + + def hybrid_search( query: str, top_k: int = 5, @@ -923,41 +1030,14 @@ def _rec(r: dict, **fields) -> None: r["score"] = (1.0 - hw) * r["score"] + hw * h _rec(r, hotness=round(h, 4), after_hotness=round(r["score"], 4)) - # Extract query-matched entities (used for co-occurrence boost AND reinforcement) - query_words = [w.lower() for w in query.split() if len(w) > 2] - query_entities: set[str] = set() - if query_words: - for word in query_words: - ent_rows = conn.execute( - "SELECT DISTINCT subject FROM triples WHERE LOWER(subject) LIKE ? " - "UNION " - "SELECT DISTINCT object FROM triples WHERE LOWER(object) LIKE ?", - (f"%{word}%", f"%{word}%"), - ).fetchall() - query_entities.update(r[0] for r in ent_rows) + # Entities the query names, used for the co-occurrence boost AND reinforcement + query_entities = _extract_query_entities(conn, query) # Co-occurrence boost: notes containing entities that co-occur with query entities # get a bounded multiplicative boost. Slots after hotness, before demotion. cooc_weight = weights.cooccurrence_boost_weight if cooc_weight > 0 and query_entities and "cooccurrence" not in ablate: - # Step 2: Find co-occurring entities and their weights - cooc_entities = {} # entity -> max co-occurrence weight - for qe in query_entities: - # Blend structural weight with accumulated search - # reinforcement (issue #60) - rows = conn.execute( - "SELECT entity_b, weight + reinforcement AS w " - "FROM entity_cooccurrence WHERE entity_a = ? " - "UNION ALL " - "SELECT entity_a, weight + reinforcement AS w " - "FROM entity_cooccurrence WHERE entity_b = ?", - (qe, qe), - ).fetchall() - for r in rows: - ent = r[0] - w = r[1] - if ent not in query_entities: # Don't boost for direct matches - cooc_entities[ent] = max(cooc_entities.get(ent, 0), w) + cooc_entities = _cooccurring_entities(conn, query_entities) if cooc_entities: # Step 3: Build note -> entities map from triples for result notes @@ -1125,18 +1205,41 @@ def _rec(r: dict, **fields) -> None: # Hebbian reinforcement: strengthen co-occurrence for entity pairs shared # between query-matched entities and result-note entities. # Fires regardless of cooccurrence_boost_weight setting. + # + # Buffered, never written here (issue #120). This used to be a write + # inside the search request, and SQLite has one writer: two overlapping + # searches serialized on the lock and the second was measured stalling + # for minutes. buffer_reinforcement queues the pairs and the write + # happens off the request path — a background drain once the buffer + # fills, `neurostack cooccurrence --flush`, or process exit — with the + # same reinforce_cooccurrence semantics. if query_entities and returned_paths: try: placeholders = ",".join("?" * len(returned_paths)) result_ent_rows = conn.execute( - f"SELECT DISTINCT subject, object FROM triples " + f"SELECT DISTINCT note_path, subject, object FROM triples " f"WHERE note_path IN ({placeholders})", returned_paths, ).fetchall() - result_entities: set[str] = set() + by_path: dict[str, set[str]] = {} for rer in result_ent_rows: - result_entities.add(rer["subject"]) - result_entities.add(rer["object"]) + ents = by_path.setdefault(rer["note_path"], set()) + ents.add(rer["subject"]) + ents.add(rer["object"]) + + # Walk the results in rank order and stop at + # MAX_RESULT_ENTITIES: the top note's entities are the ones + # whose association with the query is worth strengthening. + result_entities: list[str] = [] + seen_entities: set[str] = set() + for path in returned_paths: + for ent in sorted(by_path.get(path, ())): + if ent not in seen_entities: + seen_entities.add(ent) + result_entities.append(ent) + if len(result_entities) >= MAX_RESULT_ENTITIES: + break + del result_entities[MAX_RESULT_ENTITIES:] # Build reinforcement pairs: each query entity x each result entity reinforce_pairs = [ @@ -1146,7 +1249,7 @@ def _rec(r: dict, **fields) -> None: if qe != re ] if reinforce_pairs: - reinforce_cooccurrence(conn, reinforce_pairs) + buffer_reinforcement(conn, reinforce_pairs) except Exception: pass # Never let reinforcement disrupt search diff --git a/tests/conftest.py b/tests/conftest.py index da85918..b028453 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -97,6 +97,31 @@ def tmp_vault(tmp_path): return vault +@pytest.fixture(autouse=True) +def clear_reinforcement_buffer(): + """Isolate the per-process reinforcement buffer between tests (issue #120). + + hybrid_search buffers co-occurrence pairs in a module-level set instead of + writing them, so without this a search in one test leaves pairs that a drain + in another would write to an unrelated database. The recorded database path + is reset too, for the same reason. + """ + from neurostack import cooccurrence + + def _reset(): + thread = cooccurrence._flush_thread + if thread is not None and thread.is_alive(): + thread.join(timeout=10.0) + with cooccurrence._reinforcement_lock: + cooccurrence._reinforcement_buffer.clear() + cooccurrence._reinforcement_db_path = "" + cooccurrence._flush_thread = None + + _reset() + yield + _reset() + + @pytest.fixture def in_memory_db(): """Create an in-memory SQLite database with the NeuroStack schema.""" diff --git a/tests/test_cooccurrence.py b/tests/test_cooccurrence.py index 554cbc8..35008c5 100644 --- a/tests/test_cooccurrence.py +++ b/tests/test_cooccurrence.py @@ -2,9 +2,12 @@ from neurostack.cooccurrence import ( MAX_COOCCURRENCE_WEIGHT, + buffer_reinforcement, + flush_reinforcement, get_cooccurrence_stats, persist_cooccurrence, reinforce_cooccurrence, + reinforcement_buffer_size, upsert_cooccurrence_for_note, ) @@ -512,3 +515,191 @@ def test_upsert_keeps_reinforced_pair_when_structure_vanishes(in_memory_db): assert ("Alpha", "Ceta") not in rows assert ("Beta", "Ceta") not in rows assert rows[("Alpha", "Zeta")][0] == 1.0 + + +# --- issue #120: reinforcement is buffered, not written on the search path --- + + +def test_buffer_does_not_write(in_memory_db): + """Buffering leaves the table alone and reports what is pending.""" + conn = in_memory_db + + pending = buffer_reinforcement(conn, [("Alpha", "Beta"), ("Alpha", "Gamma")]) + + assert pending == 2 + assert reinforcement_buffer_size() == 2 + assert conn.execute( + "SELECT COUNT(*) as c FROM entity_cooccurrence" + ).fetchone()["c"] == 0 + + +def test_buffer_canonicalizes_and_dedups(in_memory_db): + """(Z, A) and (A, Z) are the same association, buffered once.""" + conn = in_memory_db + + buffer_reinforcement(conn, [("Z", "A"), ("A", "Z"), ("A", "A")]) + + assert reinforcement_buffer_size() == 1 + flush_reinforcement(conn) + row = conn.execute( + "SELECT entity_a, entity_b, reinforcement FROM entity_cooccurrence" + ).fetchone() + assert (row["entity_a"], row["entity_b"]) == ("A", "Z") + assert row["reinforcement"] == 1.0 + + +def test_flush_applies_reinforcement_maths(in_memory_db): + """Flush writes exactly what reinforce_cooccurrence would: 1.0 seed, x1.1, capped.""" + conn = in_memory_db + conn.execute( + "INSERT INTO entity_cooccurrence " + "(entity_a, entity_b, weight, reinforcement, last_seen) " + "VALUES (?, ?, ?, ?, ?)", + ("Alpha", "Beta", 3.0, 2.0, "2026-01-01"), + ) + conn.execute( + "INSERT INTO entity_cooccurrence " + "(entity_a, entity_b, weight, reinforcement, last_seen) " + "VALUES (?, ?, ?, ?, ?)", + ("Delta", "Gamma", 0.0, MAX_COOCCURRENCE_WEIGHT, "2026-01-01"), + ) + conn.commit() + + buffer_reinforcement(conn, [("Alpha", "Beta"), ("Delta", "Gamma"), ("X", "Y")]) + written = flush_reinforcement(conn) + + assert written == 3 + assert reinforcement_buffer_size() == 0 + rows = { + (r["entity_a"], r["entity_b"]): (r["weight"], r["reinforcement"]) + for r in conn.execute( + "SELECT entity_a, entity_b, weight, reinforcement " + "FROM entity_cooccurrence" + ) + } + assert abs(rows[("Alpha", "Beta")][1] - 2.2) < 1e-9 # 2.0 * 1.1 + assert rows[("Alpha", "Beta")][0] == 3.0 # structural untouched + assert rows[("Delta", "Gamma")][1] == MAX_COOCCURRENCE_WEIGHT # capped + assert rows[("X", "Y")] == (0.0, 1.0) # seeded + + +def test_flush_empty_buffer_is_noop(in_memory_db): + conn = in_memory_db + assert flush_reinforcement(conn) == 0 + + +def test_threshold_drains_in_background_without_explicit_call(tmp_path, monkeypatch): + """Crossing the threshold drains the buffer off the caller's thread.""" + import neurostack.cooccurrence as cooc_mod + from neurostack.schema import get_db + + conn = get_db(tmp_path / "drain.db") + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_FLUSH_THRESHOLD", 3) + + assert buffer_reinforcement(conn, [("A", "B"), ("A", "C")]) == 2 + assert reinforcement_buffer_size() == 2 + assert cooc_mod._flush_thread is None, "under the threshold, no writer starts" + + buffer_reinforcement(conn, [("A", "D"), ("A", "E")]) + cooc_mod._flush_thread.join(timeout=10.0) + + assert reinforcement_buffer_size() == 0 + assert conn.execute( + "SELECT COUNT(*) as c FROM entity_cooccurrence" + ).fetchone()["c"] == 4, "the whole buffer is written, not just the new pairs" + + +def test_threshold_does_not_write_on_calling_thread(tmp_path, monkeypatch): + """The drain runs on another thread — the buffering caller never writes.""" + import threading + + import neurostack.cooccurrence as cooc_mod + from neurostack.schema import get_db + + conn = get_db(tmp_path / "thread.db") + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_FLUSH_THRESHOLD", 0) + writer_threads = [] + real_reinforce = cooc_mod.reinforce_cooccurrence + + def record(*args, **kwargs): + writer_threads.append(threading.current_thread()) + return real_reinforce(*args, **kwargs) + + monkeypatch.setattr(cooc_mod, "reinforce_cooccurrence", record) + + buffer_reinforcement(conn, [("A", "B")]) + cooc_mod._flush_thread.join(timeout=10.0) + + assert writer_threads, "the drain should have written" + assert threading.current_thread() not in writer_threads + + +def test_in_memory_db_never_starts_a_writer(in_memory_db, monkeypatch): + """An in-memory database has no file another thread could open.""" + import neurostack.cooccurrence as cooc_mod + + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_FLUSH_THRESHOLD", 0) + + buffer_reinforcement(in_memory_db, [("A", "B")]) + + assert cooc_mod._flush_thread is None + assert reinforcement_buffer_size() == 1, "pairs wait for an explicit flush" + + +def test_buffer_capped(in_memory_db, monkeypatch): + """Past the ceiling, pairs are dropped rather than growing the process.""" + import neurostack.cooccurrence as cooc_mod + + conn = in_memory_db + # Threshold above the cap so nothing drains and the cap is what bites + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_BUFFER_MAX", 5) + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_FLUSH_THRESHOLD", 1000) + + buffer_reinforcement(conn, [("A", f"B{i}") for i in range(20)]) + + assert reinforcement_buffer_size() == 5 + + +def test_failing_write_does_not_raise(tmp_path, monkeypatch): + """A drain whose write blows up is swallowed, and the batch is dropped.""" + import neurostack.cooccurrence as cooc_mod + from neurostack.schema import get_db + + conn = get_db(tmp_path / "boom.db") + + def boom(*args, **kwargs): + raise RuntimeError("database is locked") + + monkeypatch.setattr(cooc_mod, "reinforce_cooccurrence", boom) + monkeypatch.setattr(cooc_mod, "REINFORCEMENT_FLUSH_THRESHOLD", 0) + + buffer_reinforcement(conn, [("A", "B")]) + cooc_mod._flush_thread.join(timeout=10.0) + + assert reinforcement_buffer_size() == 0, "a failed batch is dropped, not retried" + + +def test_atexit_flush_persists_to_file_db(tmp_path): + """The atexit hook writes the pending signal via its own connection.""" + import sqlite3 + + import neurostack.cooccurrence as cooc_mod + from neurostack.schema import get_db + + db_path = tmp_path / "atexit.db" + conn = get_db(db_path) + buffer_reinforcement(conn, [("Alpha", "Beta")]) + conn.close() # as if the request thread's connection were already gone + + cooc_mod._flush_reinforcement_at_exit() + + assert reinforcement_buffer_size() == 0 + check = sqlite3.connect(str(db_path)) + check.row_factory = sqlite3.Row + row = check.execute( + "SELECT entity_a, entity_b, reinforcement FROM entity_cooccurrence" + ).fetchone() + check.close() + assert (row["entity_a"], row["entity_b"], row["reinforcement"]) == ( + "Alpha", "Beta", 1.0, + ) diff --git a/tests/test_search.py b/tests/test_search.py index 6786098..035f3ec 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -729,8 +729,12 @@ class TestReinforcementFromSearch: """Tests for Hebbian reinforcement wired into hybrid_search.""" def test_reinforce_from_search(self, in_memory_db, monkeypatch): - """After hybrid_search, co-occurrence weights increase for entity pairs - appearing in both query-matched entities and result-note entities.""" + """hybrid_search buffers the pairs; a flush is what writes them. + + Rewritten for issue #120: search no longer writes to + entity_cooccurrence, so the same fixture now asserts the pair is + buffered by the search and lands with the same reinforcement value once + flushed.""" conn = in_memory_db emb = _fake_embedding(0.5) @@ -791,13 +795,25 @@ def test_reinforce_from_search(self, in_memory_db, monkeypatch): finally: config_mod._config = original_config - # After search, reinforcement should have created/increased the - # co-occurrence between alpha (query entity) and beta (result-note entity) + # Search itself must not have written — the whole point of issue #120 + assert conn.execute( + "SELECT COUNT(*) AS c FROM entity_cooccurrence" + ).fetchone()["c"] == 0 + + # The pair is buffered, canonical order (alpha < beta) + from neurostack.cooccurrence import ( + _reinforcement_buffer as buffered, + ) + from neurostack.cooccurrence import flush_reinforcement + assert ("alpha", "beta") in buffered + + # After a flush, reinforcement lands exactly as it used to + flush_reinforcement(conn) after = conn.execute( "SELECT weight, reinforcement FROM entity_cooccurrence " "WHERE entity_a = 'alpha' AND entity_b = 'beta'" ).fetchone() - assert after is not None, "Reinforcement should have created (alpha, beta) pair" + assert after is not None, "Flush should have created (alpha, beta) pair" initial_reinf = initial["reinforcement"] if initial else 0.0 assert after["reinforcement"] > initial_reinf, ( f"Reinforcement should have increased from {initial_reinf}" @@ -846,6 +862,308 @@ def test_reinforce_noop_no_entities_in_search(self, in_memory_db, monkeypatch): ).fetchone()["c"] assert count == 0, "No reinforcement should occur when query matches no entities" + def test_search_does_not_write_cooccurrence(self, in_memory_db, monkeypatch): + """hybrid_search issues zero writes to entity_cooccurrence (issue #120). + + SQLite has one writer, so a write inside the request serialized + overlapping searches. The pairs must land in the buffer instead, leaving + the table byte-identical.""" + conn = in_memory_db + emb = _fake_embedding(0.5) + + conn.execute( + "INSERT INTO notes (path, title, content_hash, updated_at) " + "VALUES (?, ?, ?, ?)", + ("noteA.md", "Note A", "ha", "2026-01-01"), + ) + conn.execute( + "INSERT INTO chunks (note_path, heading_path, content, " + "content_hash, position, embedding) VALUES (?, ?, ?, ?, ?, ?)", + ("noteA.md", "## Test", "alpha related content", "h_a", 0, emb), + ) + conn.execute( + "INSERT INTO triples (note_path, subject, predicate, object, triple_text) " + "VALUES (?, ?, ?, ?, ?)", + ("noteA.md", "alpha", "relates_to", "beta", "alpha relates_to beta"), + ) + # A pre-existing row so "untouched" means untouched, not merely empty + conn.execute( + "INSERT INTO entity_cooccurrence " + "(entity_a, entity_b, weight, reinforcement, last_seen) " + "VALUES (?, ?, ?, ?, ?)", + ("alpha", "beta", 4.0, 2.0, "2026-01-01"), + ) + conn.commit() + + before = conn.execute( + "SELECT entity_a, entity_b, weight, reinforcement, last_seen " + "FROM entity_cooccurrence ORDER BY entity_a, entity_b" + ).fetchall() + + import neurostack.config as config_mod + import neurostack.search as search_mod + from neurostack.cooccurrence import _reinforcement_buffer as buffered + from neurostack.cooccurrence import reinforcement_buffer_size + + monkeypatch.setattr(search_mod, "get_db", lambda path: conn) + + import numpy as np + fake_emb = np.array([0.5] * 768, dtype=np.float32) + monkeypatch.setattr( + search_mod, "get_embedding", lambda q, base_url=None: fake_emb + ) + + original_config = config_mod._config + try: + cfg = Config() + cfg.cooccurrence_boost_weight = 0.5 + config_mod._config = cfg + hybrid_search("alpha", top_k=10, embed_url="http://fake") + finally: + config_mod._config = original_config + + after = conn.execute( + "SELECT entity_a, entity_b, weight, reinforcement, last_seen " + "FROM entity_cooccurrence ORDER BY entity_a, entity_b" + ).fetchall() + assert [tuple(r) for r in after] == [tuple(r) for r in before], ( + "search must not touch entity_cooccurrence" + ) + assert reinforcement_buffer_size() > 0, "pairs should be buffered instead" + assert ("alpha", "beta") in buffered + + def test_failing_reinforcement_does_not_propagate(self, in_memory_db, monkeypatch): + """A reinforcement failure never reaches the caller of hybrid_search.""" + conn = in_memory_db + emb = _fake_embedding(0.5) + + conn.execute( + "INSERT INTO notes (path, title, content_hash, updated_at) " + "VALUES (?, ?, ?, ?)", + ("noteA.md", "Note A", "ha", "2026-01-01"), + ) + conn.execute( + "INSERT INTO chunks (note_path, heading_path, content, " + "content_hash, position, embedding) VALUES (?, ?, ?, ?, ?, ?)", + ("noteA.md", "## Test", "alpha related content", "h_a", 0, emb), + ) + conn.execute( + "INSERT INTO triples (note_path, subject, predicate, object, triple_text) " + "VALUES (?, ?, ?, ?, ?)", + ("noteA.md", "alpha", "relates_to", "beta", "alpha relates_to beta"), + ) + conn.commit() + + import neurostack.config as config_mod + import neurostack.search as search_mod + + monkeypatch.setattr(search_mod, "get_db", lambda path: conn) + + import numpy as np + fake_emb = np.array([0.5] * 768, dtype=np.float32) + monkeypatch.setattr( + search_mod, "get_embedding", lambda q, base_url=None: fake_emb + ) + + def boom(*args, **kwargs): + raise RuntimeError("buffering failed") + + monkeypatch.setattr(search_mod, "buffer_reinforcement", boom) + + original_config = config_mod._config + try: + cfg = Config() + cfg.cooccurrence_boost_weight = 0.0 + config_mod._config = cfg + results = hybrid_search("alpha", top_k=10, embed_url="http://fake") + finally: + config_mod._config = original_config + + assert len(results) > 0, "search must still return results" + + +class _CountingConn: + """Wraps a connection, counting execute() calls. + + sqlite3.Connection allows no attribute assignment, so the count is taken by + delegation rather than by monkeypatching the connection. + """ + + def __init__(self, conn): + self._conn = conn + self.executes = 0 + + def execute(self, *args, **kwargs): + self.executes += 1 + return self._conn.execute(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._conn, name) + + +class TestQueryEntityExtraction: + """Tests for _extract_query_entities (issue #120).""" + + def _seed(self, conn, entities, note="n.md", predicate="relates_to"): + conn.execute( + "INSERT OR IGNORE INTO notes (path, title, content_hash, updated_at) " + "VALUES (?, ?, ?, ?)", + (note, note, f"h_{note}", "2026-01-01"), + ) + for i, ent in enumerate(entities): + conn.execute( + "INSERT INTO triples (note_path, subject, predicate, object, " + "triple_text) VALUES (?, ?, ?, ?, ?)", + (note, ent, predicate, f"obj_{note}_{i}", f"{ent} {predicate}"), + ) + conn.commit() + + def test_interior_substring_no_longer_matches(self, in_memory_db): + """The LIKE '%word%' blowup is gone: matching is anchored at a word start.""" + from neurostack.search import _extract_query_entities + + conn = in_memory_db + self._seed(conn, [ + # word-anchored: these are what the query means + "agent registry", "Azure Foundry agent", "agents", "agentic flow", + "azure-foundry-agent", + # interior substrings: the old rule pulled all of these in + "subagent_runner", "useragentstring", "reagent batch", "myagentless", + ]) + + found = _extract_query_entities(conn, "agent") + assert found == { + "agent registry", "Azure Foundry agent", "agents", "agentic flow", + "azure-foundry-agent", + } + + # The old rule for comparison: every entity containing the substring + legacy = { + r[0] for r in conn.execute( + "SELECT DISTINCT subject FROM triples " + "WHERE LOWER(subject) LIKE '%agent%'" + ).fetchall() + } + assert len(legacy) > len(found) + + def test_short_words_skipped(self, in_memory_db): + from neurostack.search import _extract_query_entities + + conn = in_memory_db + self._seed(conn, ["ai platform", "azure hosting"]) + assert _extract_query_entities(conn, "ai") == set() + + def test_cap_keeps_least_frequent_entities(self, in_memory_db): + """Over the cap, the selective entities survive and the common ones do not.""" + from neurostack.search import MAX_QUERY_ENTITIES, _extract_query_entities + + conn = in_memory_db + rare = [f"azure rare {i}" for i in range(MAX_QUERY_ENTITIES + 20)] + self._seed(conn, rare) + # "azure everywhere" fills far more triple slots than any rare entity + self._seed(conn, ["azure everywhere"] * 40, note="hot.md") + + found = _extract_query_entities(conn, "azure") + assert len(found) == MAX_QUERY_ENTITIES + assert "azure everywhere" not in found, ( + "a term appearing everywhere carries no ranking signal" + ) + + def test_extraction_is_one_query_regardless_of_word_count(self, in_memory_db): + """Query count does not scale with the number of query words.""" + from neurostack.search import _extract_query_entities + + conn = in_memory_db + self._seed(conn, ["azure foundry", "knowledge base", "agent registry"]) + + one = _CountingConn(conn) + _extract_query_entities(one, "azure") + many = _CountingConn(conn) + _extract_query_entities(many, "azure foundry knowledge agent registry base") + assert one.executes == 1 + assert many.executes == 1 + + def test_wildcards_in_query_word_are_literal(self, in_memory_db): + """A '%' in a query word must not match everything.""" + from neurostack.search import _extract_query_entities + + conn = in_memory_db + self._seed(conn, ["azure hosting", "100% coverage"]) + assert _extract_query_entities(conn, "%cov") == set() + + +class TestCooccurringEntitiesLookup: + """Tests for _cooccurring_entities (issue #120).""" + + def _reference_impl(self, conn, query_entities): + """The per-entity loop this replaced, kept as the regression oracle.""" + cooc_entities = {} + for qe in query_entities: + rows = conn.execute( + "SELECT entity_b, weight + reinforcement AS w " + "FROM entity_cooccurrence WHERE entity_a = ? " + "UNION ALL " + "SELECT entity_a, weight + reinforcement AS w " + "FROM entity_cooccurrence WHERE entity_b = ?", + (qe, qe), + ).fetchall() + for r in rows: + ent, w = r[0], r[1] + if ent not in query_entities: + cooc_entities[ent] = max(cooc_entities.get(ent, 0), w) + return cooc_entities + + def _seed_pairs(self, conn, pairs): + for a, b, weight, reinforcement in pairs: + conn.execute( + "INSERT INTO entity_cooccurrence " + "(entity_a, entity_b, weight, reinforcement, last_seen) " + "VALUES (?, ?, ?, ?, ?)", + (a, b, weight, reinforcement, "2026-01-01"), + ) + conn.commit() + + def test_result_matches_per_entity_loop(self, in_memory_db): + """Same entity -> weight mapping as the loop it replaced.""" + from neurostack.search import _cooccurring_entities + + conn = in_memory_db + self._seed_pairs(conn, [ + ("alpha", "beta", 3.0, 0.0), # query entity on side a + ("delta", "alpha", 1.0, 2.0), # query entity on side b, blended + ("alpha", "gamma", 5.0, 0.0), # both query entities -> excluded + ("beta", "gamma", 9.0, 0.0), # both sides non-query... via gamma + ("epsilon", "zeta", 7.0, 0.0), # unrelated + ("gamma", "beta2", 2.0, 0.5), # gamma is a query entity + ]) + query_entities = {"alpha", "gamma"} + + got = _cooccurring_entities(conn, query_entities) + assert got == self._reference_impl(conn, query_entities) + # spelled out, so the oracle cannot drift silently + assert got == {"beta": 9.0, "delta": 3.0, "beta2": 2.5} + + def test_lookup_is_chunk_bounded_not_per_entity(self, in_memory_db): + """One statement per SQL_PARAM_CHUNK entities, not one (or two) each.""" + from neurostack.cooccurrence import SQL_PARAM_CHUNK + from neurostack.search import _cooccurring_entities + + conn = in_memory_db + self._seed_pairs(conn, [("q0", "target", 1.0, 0.0)]) + + entities = {f"q{i}" for i in range(SQL_PARAM_CHUNK + 10)} + counting = _CountingConn(conn) + got = _cooccurring_entities(counting, entities) + assert counting.executes == 2, "expected one statement per chunk" + assert got == {"target": 1.0} + + def test_no_query_entities_issues_no_query(self, in_memory_db): + from neurostack.search import _cooccurring_entities + + counting = _CountingConn(in_memory_db) + assert _cooccurring_entities(counting, set()) == {} + assert counting.executes == 0 + class TestLinkSectionHelpers: """Unit tests for the link-section detection helpers (issue #41)."""