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