diff --git a/better_memory/services/relevant.py b/better_memory/services/relevant.py index 221b612..5ee7335 100644 --- a/better_memory/services/relevant.py +++ b/better_memory/services/relevant.py @@ -105,6 +105,7 @@ def _vec_qualifiers( id_col: str, query_vector: list[float] | None, vec_floor: float, + candidate_ids: list[str] | None = None, ) -> dict[str, int]: """id -> vec rank for rows within the cosine floor. @@ -117,20 +118,37 @@ def _vec_qualifiers( L2 distance, NOT squared L2. The floor comparison below is against the plain (unsquared) distance only; the squared-distance branch that a defensive dual-check would need is dead code and has been omitted. + + ``candidate_ids`` scopes the result to the caller's candidate set, + matching ``ReflectionService._vec_ranks`` (reflection.py) -- sqlite-vec + kNN accepts only ``embedding MATCH ? AND k = ?`` (no extra predicates), + so we fetch top-k and filter in Python. ``k`` scales with + ``max(len(candidate_ids), 50)`` so the window is dense over what the + caller actually ranks, instead of being drowned by neighbours from + other projects or retired records (whose vectors persist -- retire + only UPDATEs status). ``None`` preserves the global-scope behaviour + for callers (e.g. ``SqliteBackend.relevance_ranks``) that intentionally + want the unfiltered rank map; ``[]`` short-circuits to ``{}``. """ if conn is None or query_vector is None: return {} + if candidate_ids is not None and not candidate_ids: + return {} max_dist = math.sqrt(2.0 * (1.0 - vec_floor)) + k = max(len(candidate_ids), 50) if candidate_ids is not None else 50 try: rows = conn.execute( f"SELECT {id_col}, distance FROM {table} " f"WHERE embedding MATCH ? AND k = ? ORDER BY distance", - (sqlite_vec.serialize_float32(query_vector), 50), + (sqlite_vec.serialize_float32(query_vector), k), ).fetchall() except sqlite3.OperationalError: return {} + wanted: set[str] | None = set(candidate_ids) if candidate_ids is not None else None out: dict[str, int] = {} for row in rows: + if wanted is not None and row[0] not in wanted: + continue if float(row[1]) <= max_dist: out[row[0]] = len(out) return out @@ -213,8 +231,26 @@ def retrieve_relevant( bm = _bm25_qualifiers(conn, query) qvec = sync_embedder.embed_text(query) if sync_embedder is not None else None - vec_r = _vec_qualifiers(conn, "reflection_embeddings", "reflection_id", qvec, vec_floor) - vec_s = _vec_qualifiers(conn, "semantic_embeddings", "memory_id", qvec, vec_floor) + # Collect the actual candidate id set before the vec query so kNN is + # restricted to it (issue #104): the vec tables are global across + # projects and retained/retired rows, so an unbounded k=50 kNN can be + # entirely filled by neighbours the caller will never rank -- starving + # the evidence gate for the caller's own reflections/semantics. + refl_bucket_order = ["do", "dont"] + (["neutral"] if include_neutral else []) + refl_candidate_ids = [ + str(r.get("id")) + for bucket in refl_bucket_order + for r in (buckets.get(bucket) or []) + ] + sem_candidate_ids = [str(getattr(s, "id", "")) for s in (semantic or [])] + vec_r = _vec_qualifiers( + conn, "reflection_embeddings", "reflection_id", qvec, vec_floor, + refl_candidate_ids, + ) + vec_s = _vec_qualifiers( + conn, "semantic_embeddings", "memory_id", qvec, vec_floor, + sem_candidate_ids, + ) keywords = extract_keywords(query) # fallback evidence only fts_unavailable = conn is None @@ -255,8 +291,7 @@ def retrieve_relevant( rank_map: dict[tuple[str, str], int] = raw_rank_map or {} refl_candidates: list[dict] = [] - order = ["do", "dont"] + (["neutral"] if include_neutral else []) - for bucket in order: + for bucket in refl_bucket_order: for r in buckets.get(bucket, []) or []: r_id = str(r.get("id")) title = str(r.get("title") or "") diff --git a/tests/services/test_relevant.py b/tests/services/test_relevant.py index d7195a2..c22a17c 100644 --- a/tests/services/test_relevant.py +++ b/tests/services/test_relevant.py @@ -144,6 +144,33 @@ def test_vec_below_floor_does_not_qualify(self, conn): ) assert out == [] + def test_orphan_vector_row_does_not_leak_via_vec_gate(self, conn): + # Issue #104: the vec table is global and its rows persist across + # retire (retire only UPDATEs status; the vector stays). An + # "orphan" vector -- one whose metadata row was retired or deleted + # so backend.retrieve does not surface it as a candidate -- must + # not qualify anything through the vec leg. Prior to the fix, + # _vec_qualifiers issued an unfiltered kNN, so the orphan showed + # up in vec_r but happened to be harmless only because the RRF + # loop keys on the caller's candidate list -- this test locks in + # that the candidate scoping is now enforced at the vec-query + # layer too, so orphan ids never appear in vec_r at all. + from better_memory.services.relevant import _vec_qualifiers + emb = DirectedEmbedder("orphan trigger") + # Insert a vec row for a reflection id that has NO metadata row. + _embed_reflection(conn, "r-orphan", emb._vec("orphan trigger")) + qvec = emb._vec("orphan trigger") + # candidate_ids empty -> short-circuit + assert _vec_qualifiers( + conn, "reflection_embeddings", "reflection_id", + qvec, 0.55, [], + ) == {} + # candidate_ids restricts what returns from the kNN scan + assert _vec_qualifiers( + conn, "reflection_embeddings", "reflection_id", + qvec, 0.55, ["r-different"], + ) == {} + class TestWilsonRanking: def test_wilson_ranks_among_qualifiers(self, conn):