@pentatonic-ai/ai-agent-sdk 0.10.7 → 0.10.8

package/packages/memory-engine-v2/scripts/fusion_drive_fuse.py

@@ -0,0 +1,264 @@
+#!/usr/bin/env python3
+"""Fusion Drive — fusion pass (detect duplicate nodes, fuse into a master).
+
+Detects duplicate/near-duplicate entities (and exact-triple duplicate facts)
+within an arena and merges each set into a single master, repointing facts +
+relationships, unioning aliases/provenance, and writing reversible audit rows
+(entity_merges / fact_merges).
+
+Detection here is the TRACTABLE, no-LLM tier:
+  - entities: exact normalized-name / alias-overlap dupes, AND cross-run
+    shared-provenance dupes (two same-type entities citing the same event
+    where one is junk-leaning — the 7B numeric-ID vs new-teacher real-name
+    case). Embedding-band + LLM adjudication detection stays in
+    entity_resolution_v2.py (#82); this pass consumes its proposals too if
+    handed a --proposals file.
+  - facts: exact (subject, predicate, object) triples.
+Master selection uses fusion_drive.canonical (directory-anchored scoring),
+NOT richest-row-wins.
+
+Safety: arena-scoped (required), DRY-RUN default, --apply to merge, each merge
+its own transaction, every deprecated row recoverable from the audit table,
+restricted disclosure never auto-merged.
+
+Usage:
+  fusion_drive_fuse.py --arena 'X'           # dry-run: list merge proposals
+  fusion_drive_fuse.py --arena 'X' --apply   # execute merges (reversible)
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import uuid
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "fusion_drive"))
+import canonical as C  # noqa: E402
+from merge import build_entity_merge_plan, build_fact_merge_plan  # noqa: E402
+
+try:
+    import psycopg
+    from psycopg.rows import dict_row
+except ModuleNotFoundError:
+    print("psycopg required", file=sys.stderr)
+    raise
+
+
+def _norm(s: str) -> str:
+    return " ".join(s.lower().split())
+
+
+def _entity_dup_sets(cur, arena: str) -> list[list[dict]]:
+    """Group same-(type) entities that are exact normalized-name dupes OR
+    share a provenance event with a junk-leaning twin. Returns groups of >=2."""
+    cur.execute(
+        """SELECT id, entity_type, canonical_name, aliases, provenance_event_ids, disclosure_class
+           FROM entities WHERE arena = %s AND disclosure_class <> 'restricted'""",
+        (arena,),
+    )
+    ents = cur.fetchall()
+    groups: dict[tuple, list[dict]] = {}
+    # 1. exact normalized-name within (type)
+    for e in ents:
+        key = (e["entity_type"], _norm(e["canonical_name"]))
+        groups.setdefault(key, []).append(e)
+    exact = [g for g in groups.values() if len(g) > 1]
+
+    # 2. cross-run shared-provenance: same type + same event in provenance,
+    #    where some members are junk-leaning (looks-like-id) — catches
+    #    name-divergent dupes like "1716801984" vs "Katie Cooper" that never
+    #    block on name.
+    #
+    #    OVER-MERGE GUARD: a single event can legitimately mention several
+    #    distinct same-type entities (an email naming Alice, Bob, AND a
+    #    numeric-ID node). Merging the whole co-occurrence group would
+    #    conflate Alice and Bob. So the no-LLM tier ONLY proposes when the
+    #    group has EXACTLY ONE non-junk member: we fold the junk node(s) into
+    #    that unambiguous real master. Groups with 0 or >=2 non-junk members
+    #    are ambiguous and deferred to the LLM-adjudicated tier
+    #    (entity_resolution_v2.py) rather than auto-merged.
+    by_event_type: dict[tuple, list[dict]] = {}
+    for e in ents:
+        for ev in (e["provenance_event_ids"] or []):
+            by_event_type.setdefault((e["entity_type"], ev), []).append(e)
+    cross = []
+    seen_ids: set[str] = set()
+    for members in by_event_type.values():
+        if len(members) < 2:
+            continue
+        junk = [m for m in members if C.looks_like_id(m["canonical_name"])]
+        non_junk = [m for m in members if not C.looks_like_id(m["canonical_name"])]
+        if not junk or len(non_junk) != 1:
+            continue  # need junk to clean AND exactly one unambiguous master
+        group = non_junk + junk
+        ids = tuple(sorted(m["id"] for m in group))
+        if ids in seen_ids:
+            continue
+        seen_ids.add(ids)
+        cross.append(group)
+    return exact + cross
+
+
+def _candidates(group: list[dict]) -> list[C.CanonicalCandidate]:
+    return [
+        C.CanonicalCandidate(
+            entity_id=e["id"],
+            canonical_name=e["canonical_name"],
+            n_provenance=len(e["provenance_event_ids"] or []),
+            aliases=e["aliases"] or [],
+            # in_directory / grounded / from_current_teacher would be resolved
+            # from an authority table + provenance content + trace llm_model;
+            # left False here (no-LLM tier) so scoring leans on grounding-by-
+            # corroboration + the ID/bare penalties. Wire authority in a
+            # follow-up — the scoring already supports it.
+        )
+        for e in group
+    ]
+
+
+def _touching(cur, arena: str, loser_ids: list[str]) -> tuple[list[dict], list[dict]]:
+    cur.execute(
+        """SELECT id, subject_entity_id, object_entity_id FROM facts
+           WHERE arena = %s AND (subject_entity_id = ANY(%s) OR object_entity_id = ANY(%s))""",
+        (arena, loser_ids, loser_ids),
+    )
+    facts = cur.fetchall()
+    cur.execute(
+        """SELECT id, from_entity_id, to_entity_id, relationship_type, weight, provenance_event_ids
+           FROM relationships WHERE arena = %s
+           AND (from_entity_id = ANY(%s) OR to_entity_id = ANY(%s)
+                OR from_entity_id = ANY(%s) OR to_entity_id = ANY(%s))""",
+        (arena, loser_ids, loser_ids, loser_ids, loser_ids),
+    )
+    rels = cur.fetchall()
+    return facts, rels
+
+
+def _execute_entity_plan(cur, plan) -> None:
+    # master row
+    cur.execute("UPDATE entities SET aliases = %s, provenance_event_ids = %s, last_seen = NOW() WHERE id = %s",
+                (plan.master_aliases, plan.master_provenance, plan.master_id))
+    for fid in plan.fact_subject_repoints:
+        cur.execute("UPDATE facts SET subject_entity_id = %s WHERE id = %s", (plan.master_id, fid))
+    for fid in plan.fact_object_repoints:
+        cur.execute("UPDATE facts SET object_entity_id = %s WHERE id = %s", (plan.master_id, fid))
+    for rid in plan.rel_endpoint_repoints:
+        cur.execute(
+            """UPDATE relationships SET
+                 from_entity_id = CASE WHEN from_entity_id = ANY(%s) THEN %s ELSE from_entity_id END,
+                 to_entity_id   = CASE WHEN to_entity_id   = ANY(%s) THEN %s ELSE to_entity_id END
+               WHERE id = %s""",
+            (plan.deprecated_entity_ids, plan.master_id,
+             plan.deprecated_entity_ids, plan.master_id, rid),
+        )
+    for col in plan.rel_collisions:
+        cur.execute("UPDATE relationships SET weight = %s, provenance_event_ids = %s WHERE id = %s",
+                    (col["summed_weight"], col["provenance"], col["keep"]))
+        cur.execute("DELETE FROM relationships WHERE id = %s", (col["drop"],))
+    for a in plan.audit_rows:
+        cur.execute(
+            """INSERT INTO entity_merges (id, arena, canonical_id, deprecated_id,
+                 deprecated_canonical_name, deprecated_aliases, merge_signal,
+                 facts_repointed, rollback_payload)
+               VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s::jsonb)""",
+            ("em_" + uuid.uuid4().hex[:20], a["arena"], a["canonical_id"], a["deprecated_id"],
+             a["deprecated_canonical_name"], a["deprecated_aliases"], a["merge_signal"],
+             len(plan.fact_subject_repoints) + len(plan.fact_object_repoints),
+             json.dumps(a["rollback_payload"], default=str)),
+        )
+    cur.execute("DELETE FROM entities WHERE id = ANY(%s)", (plan.deprecated_entity_ids,))
+
+
+def _dedup_master_facts(cur, arena: str, master_id: str) -> int:
+    """After repointing facts onto the master, the master can hold several
+    facts with the same (subject, predicate, object) but different statements
+    (fact id is content_id(arena, statement), so they didn't collapse on
+    insert). Fuse each such triple-group via build_fact_merge_plan: keep the
+    best, union provenance, delete dups with a fact_merges receipt."""
+    cur.execute(
+        """SELECT id, predicate, object_entity_id, statement, confidence, provenance_event_ids
+           FROM facts
+           WHERE arena = %s AND (subject_entity_id = %s OR object_entity_id = %s)""",
+        (arena, master_id, master_id),
+    )
+    rows = cur.fetchall()
+    groups: dict[tuple, list[dict]] = {}
+    for r in rows:
+        # group key uses the master as the subject anchor + predicate + object
+        groups.setdefault((master_id, r["predicate"], r["object_entity_id"]), []).append(r)
+    deduped = 0
+    for dup in groups.values():
+        plan = build_fact_merge_plan(arena=arena, dup_facts=dup)
+        if not plan:
+            continue
+        cur.execute("UPDATE facts SET provenance_event_ids = %s WHERE id = %s",
+                    (plan["master_provenance"], plan["master_id"]))
+        for a in plan["audit_rows"]:
+            cur.execute(
+                """INSERT INTO fact_merges (id, arena, canonical_id, deprecated_id,
+                     deprecated_statement, merge_signal, provenance_unioned, rollback_payload)
+                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s::jsonb)""",
+                ("fm_" + uuid.uuid4().hex[:20], a["arena"], a["canonical_id"], a["deprecated_id"],
+                 a["deprecated_statement"], a["merge_signal"], a["provenance_unioned"],
+                 json.dumps(a["rollback_payload"], default=str)),
+            )
+        cur.execute("DELETE FROM facts WHERE id = ANY(%s)", (plan["deprecated_ids"],))
+        deduped += len(plan["deprecated_ids"])
+    return deduped
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--arena", required=True)
+    ap.add_argument("--pg-dsn", default=os.environ.get("PG_DSN"))
+    ap.add_argument("--apply", action="store_true", help="execute merges (default: dry-run)")
+    args = ap.parse_args()
+    if not args.pg_dsn:
+        print("PG_DSN required", file=sys.stderr)
+        return 2
+
+    proposals = 0
+    merged = 0
+    with psycopg.connect(args.pg_dsn, row_factory=dict_row) as conn:
+        with conn.cursor() as cur:
+            groups = _entity_dup_sets(cur, args.arena)
+            for group in groups:
+                master_c, losers_c = C.pick_master(_candidates(group))
+                loser_ids = [l.entity_id for l in losers_c]
+                if not loser_ids:
+                    continue
+                proposals += 1
+                by_id = {e["id"]: e for e in group}
+                master = by_id[master_c.entity_id]
+                losers = [by_id[i] for i in loser_ids]
+                facts, rels = _touching(cur, args.arena, loser_ids)
+                plan = build_entity_merge_plan(
+                    arena=args.arena, master=master, losers=losers, facts=facts, relationships=rels)
+                print(f"  MERGE → master '{master['canonical_name']}' ({master['id']}) "
+                      f"absorbs {[l['canonical_name'] for l in losers]} "
+                      f"[facts:{len(plan.fact_subject_repoints)+len(plan.fact_object_repoints)} "
+                      f"rels:{len(plan.rel_endpoint_repoints)} collisions:{len(plan.rel_collisions)}]")
+                if args.apply:
+                    _execute_entity_plan(cur, plan)
+                    _dedup_master_facts(cur, args.arena, master["id"])
+                    merged += len(loser_ids)
+                    conn.commit()  # per-merge: a bad merge can't roll back the good ones, and locks stay short
+            run_id = "fdr_" + uuid.uuid4().hex[:20]
+            cur.execute(
+                """INSERT INTO fusion_drive_runs (id, arena, pass_kind, mode, scanned, changed, detail, finished_at)
+                   VALUES (%s,%s,'fusion',%s,%s,%s,%s::jsonb,NOW())""",
+                (run_id, args.arena, "apply" if args.apply else "dry_run",
+                 proposals, merged, json.dumps({"proposals": proposals, "merged": merged})),
+            )
+        conn.commit()
+
+    label = "APPLY (merged, reversible via entity_merges)" if args.apply else "DRY-RUN"
+    print(f"[fusion-drive:fuse] {label} arena={args.arena}: {proposals} proposal(s), {merged} entities merged")
+    print(f"  ledger: {run_id}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())