diff --git a/kgmd/resolve.py b/kgmd/resolve.py index 8ba1520..9a4f1c4 100644 --- a/kgmd/resolve.py +++ b/kgmd/resolve.py @@ -112,8 +112,7 @@ def run_resolution(conn, config: dict, corpus_dir: Path | None = None) -> dict: conn, merge_ids ) - _merge_entities(conn, survivor_id, merge_ids[1:], canonical) - total_merges += len(merge_ids) - 1 + total_merges += _merge_entities(conn, survivor_id, merge_ids[1:], canonical) _finalize_run(conn, run_id, total_merges) return {"merges": total_merges} @@ -238,8 +237,15 @@ def _most_frequent_surface(conn, entity_ids: list[int]) -> str: return row["surface_form"] if row else "" -def _merge_entities(conn, survivor_id: int, drop_ids: list[int], canonical_name: str) -> None: - """Merge dropped entities into the survivor.""" +def _merge_entities(conn, survivor_id: int, drop_ids: list[int], canonical_name: str) -> int: + """Merge dropped entities into the survivor. + + Returns the number of entities actually merged away: 0 when the merge is + skipped because the survivor row is already gone, and otherwise only the + drop_ids that still had a row to merge. Both cases arise because the + mentions list is read once per run, so overlapping clusters can name a + survivor or a drop that an earlier cluster already deleted. + """ # The `mentions` list is fetched once at the start of run_resolution; later # clusters in the same run may name a survivor_id that an earlier cluster # already deleted. Skip silently in that case — the entity was already @@ -248,14 +254,16 @@ def _merge_entities(conn, survivor_id: int, drop_ids: list[int], canonical_name: "SELECT attributes FROM entities WHERE id = ?", (survivor_id,) ).fetchone() if survivor_row is None: - return + return 0 # Merge attributes before deleting drops survivor_attrs = json.loads(survivor_row["attributes"]) + merged = 0 for did in drop_ids: row = conn.execute("SELECT attributes FROM entities WHERE id = ?", (did,)).fetchone() if row: other = json.loads(row["attributes"]) survivor_attrs.update(other) + merged += 1 conn.execute( "UPDATE entities SET attributes = ? WHERE id = ?", @@ -317,6 +325,7 @@ def _merge_entities(conn, survivor_id: int, drop_ids: list[int], canonical_name: ) conn.commit() + return merged def _finalize_run(conn, run_id: int, merges: int) -> None: diff --git a/tests/test_resolve.py b/tests/test_resolve.py index 33aa0e5..56c54fa 100644 --- a/tests/test_resolve.py +++ b/tests/test_resolve.py @@ -198,3 +198,120 @@ def test_merge_entities_relation_unique_collision(initialized_corpus): assert survivor_rel == 1 conn.close() + + +def test_resolution_merge_count_excludes_skipped_cluster(initialized_corpus): + """The reported merge count must not include a cluster that was skipped. + + The mentions list is read once at the start of the run, so a later cluster + can name a survivor_id an earlier cluster already deleted. _merge_entities + skips that cluster; the count run_resolution persists and returns has to + skip it too, or resolution_runs.merges overstates what happened. + """ + from kgmd.config import load_config + from kgmd.db import get_connection + + db_path = initialized_corpus / ".kgmd" / "graph.db" + conn = get_connection(db_path) + config = load_config(initialized_corpus) + config["resolution"]["llm_verify_clusters"] = False + now = datetime.now(timezone.utc).isoformat() + + _seed_run_doc_chunk(conn, now) + + # Three entities of the same type. Entity 2 carries two mentions, one + # matching entity 1 and one matching entity 3, so union-find produces two + # separate clusters: {1, 2} and {2, 3}. + conn.execute(SQL_INSERT_ENTITY, ("Brian Anderson", "Person", now, now)) + conn.execute(SQL_INSERT_ENTITY, ("B. Anderson", "Person", now, now)) + conn.execute(SQL_INSERT_ENTITY, ("Bri Anderson", "Person", now, now)) + + dim = 384 + vec_a = [0.0] * dim + vec_a[0] = 1.0 + vec_b = [0.0] * dim + vec_b[1] = 1.0 # orthogonal to vec_a -> cosine similarity 0 + + mentions = [ + (1, "Brian Anderson", vec_a), + (2, "B. Anderson", vec_a), + (2, "B Anderson", vec_b), + (3, "Bri Anderson", vec_b), + ] + for mention_id, (entity_id, surface, vec) in enumerate(mentions, start=1): + conn.execute(SQL_INSERT_MENTION, (entity_id, surface)) + conn.execute( + "INSERT INTO vec_entity_mentions (mention_id, embedding) VALUES (?, ?)", + (mention_id, struct.pack(f"{dim}f", *vec)), + ) + conn.commit() + + stats = run_resolution(conn, config) + + # Only the first cluster merges: entity 2 folds into entity 1. The second + # cluster names survivor_id 2, which no longer exists, so it is skipped and + # entity 3 survives. + remaining = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0] + assert remaining == 2 + assert stats["merges"] == 1 + + persisted = conn.execute( + "SELECT merges FROM resolution_runs ORDER BY id DESC LIMIT 1" + ).fetchone()[0] + assert persisted == 1 + + conn.close() + + +def test_resolution_merge_count_excludes_already_deleted_drop(initialized_corpus): + """A drop entity an earlier cluster already deleted must not be counted. + + The mirror of the stale-survivor case: here the survivor is still present, + so the merge runs, but one of its drop_ids was deleted by an earlier + cluster. Re-pointing and deleting that id are no-ops, so it must not be + counted as a merge. + """ + from kgmd.config import load_config + from kgmd.db import get_connection + + db_path = initialized_corpus / ".kgmd" / "graph.db" + conn = get_connection(db_path) + config = load_config(initialized_corpus) + config["resolution"]["llm_verify_clusters"] = False + now = datetime.now(timezone.utc).isoformat() + + _seed_run_doc_chunk(conn, now) + + # Entity 3 is shared between two clusters: {2, 3} and {1, 3}. The first + # deletes entity 3 as a drop; the second still names it as a drop under a + # survivor (entity 1) that is very much alive. + conn.execute(SQL_INSERT_ENTITY, ("Brian Anderson", "Person", now, now)) + conn.execute(SQL_INSERT_ENTITY, ("Bri Anderson", "Person", now, now)) + conn.execute(SQL_INSERT_ENTITY, ("B. Anderson", "Person", now, now)) + + dim = 384 + vec_a = [0.0] * dim + vec_a[0] = 1.0 + vec_b = [0.0] * dim + vec_b[1] = 1.0 + + mentions = [ + (2, "Bri Anderson", vec_a), + (3, "B. Anderson", vec_a), + (3, "B Anderson", vec_b), + (1, "Brian Anderson", vec_b), + ] + for mention_id, (entity_id, surface, vec) in enumerate(mentions, start=1): + conn.execute(SQL_INSERT_MENTION, (entity_id, surface)) + conn.execute( + "INSERT INTO vec_entity_mentions (mention_id, embedding) VALUES (?, ?)", + (mention_id, struct.pack(f"{dim}f", *vec)), + ) + conn.commit() + + stats = run_resolution(conn, config) + + remaining = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0] + assert stats["merges"] == 3 - remaining + + conn.close()