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

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

@@ -0,0 +1,434 @@
+#!/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
+from adjudicate import adjudicate_entities, adjudicate_facts  # noqa: E402
+
+try:
+    import httpx
+    import psycopg
+    from psycopg.rows import dict_row
+except ModuleNotFoundError:
+    print("psycopg (+ httpx for --llm-endpoint) required", file=sys.stderr)
+    raise
+
+# Current distiller served-model — the in-VPC LLM used for adjudication
+# (no egress; same model that extracted this content). Override with --model.
+DEFAULT_MODEL = "qwen3.6-27b-fp8"
+
+
+def _distiller_post_fn(endpoint: str, model: str):
+    """Build a post_fn(messages)->str hitting the distiller's OpenAI
+    /v1/chat/completions (temperature 0, thinking off — same shape the worker
+    uses). In-VPC: memory content never leaves the network. Raises on failure
+    so the adjudicator fails closed (treats it as 'unsure', never merges)."""
+    def post(messages):
+        r = httpx.post(endpoint, timeout=60, json={
+            "model": model, "messages": messages, "temperature": 0.0,
+            "max_tokens": 120, "chat_template_kwargs": {"enable_thinking": False},
+        })
+        r.raise_for_status()
+        return r.json()["choices"][0]["message"]["content"]
+    return post
+
+
+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 _authority_signals(cur, arena: str, entity_ids: list[str], current_model: str) -> dict:
+    """Batch-resolve the canonical-scoring authority signals (#3) for a set of
+    entities, from data that actually exists:
+      - grounded: the entity's canonical_name appears verbatim in the content
+        of at least one of its provenance events (not a hallucinated name).
+      - from_current_teacher: at least one provenance event was distilled by
+        the CURRENT teacher (distillation_traces.llm_model = current_model) —
+        prefer the newer teacher's rendering over a superseded one.
+    (in_directory is left False — there's no authoritative directory/contacts
+    table in the schema yet; that's a separate data-source decision, noted in
+    the RFC. The scorer already supports it for when one lands.)
+    Returns {entity_id: {"grounded": bool, "from_current_teacher": bool}}."""
+    out = {eid: {"grounded": False, "from_current_teacher": False} for eid in entity_ids}
+    if not entity_ids:
+        return out
+    cur.execute(
+        """SELECT e.id, e.canonical_name,
+                  EXISTS (SELECT 1 FROM events ev
+                          WHERE ev.id = ANY(e.provenance_event_ids)
+                            AND position(e.canonical_name in ev.content) > 0) AS grounded,
+                  EXISTS (SELECT 1 FROM distillation_traces t
+                          WHERE t.event_id = ANY(e.provenance_event_ids)
+                            AND t.llm_model = %s) AS cur_teacher
+           FROM entities e WHERE e.arena = %s AND e.id = ANY(%s)""",
+        (current_model, arena, entity_ids),
+    )
+    for eid, name, grounded, cur_teacher in cur.fetchall():
+        # A numeric-ID-as-person name (e.g. "1716801984") substring-matches any
+        # stray digit-run in content (epochs, order/invoice numbers) → would
+        # falsely mark junk "grounded" and BOOST its authority. Never credit a
+        # looks-like-id name as grounded. (#96 review §2)
+        grounded = bool(grounded) and not C.looks_like_id(name)
+        out[eid] = {"grounded": grounded, "from_current_teacher": bool(cur_teacher)}
+    return out
+
+
+def _candidates(group: list[dict], signals: dict | None = None) -> list[C.CanonicalCandidate]:
+    signals = signals or {}
+    return [
+        C.CanonicalCandidate(
+            entity_id=e["id"],
+            canonical_name=e["canonical_name"],
+            n_provenance=len(e["provenance_event_ids"] or []),
+            aliases=e["aliases"] or [],
+            grounded=signals.get(e["id"], {}).get("grounded", False),
+            from_current_teacher=signals.get(e["id"], {}).get("from_current_teacher", False),
+        )
+        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 _ambiguous_cross_run(cur, arena: str) -> list[dict]:
+    """Co-occurrence groups the deterministic tier PUNTED on: a junk-leaning
+    node sharing an event with MULTIPLE non-junk candidates (so which real
+    entity it belongs to is ambiguous). Returns [{junk, candidates:[...]}] for
+    the LLM tier to adjudicate. Restricted disclosure excluded."""
+    cur.execute(
+        """SELECT id, entity_type, canonical_name, aliases, provenance_event_ids
+           FROM entities WHERE arena = %s AND disclosure_class <> 'restricted'""",
+        (arena,),
+    )
+    ents = cur.fetchall()
+    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)
+    out, seen = [], set()
+    for members in by_event_type.values():
+        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) < 2:
+            continue
+        for j in junk:
+            key = (j["id"], tuple(sorted(c["id"] for c in non_junk)))
+            if key in seen:
+                continue
+            seen.add(key)
+            out.append({"junk": j, "candidates": non_junk})
+    return out
+
+
+def _entity_context(cur, arena: str, eid: str) -> list[str]:
+    cur.execute(
+        "SELECT statement FROM facts WHERE arena=%s AND (subject_entity_id=%s OR object_entity_id=%s) LIMIT 5",
+        (arena, eid, eid))
+    return [r["statement"] for r in cur.fetchall()]
+
+
+def _semantic_fact_groups(cur, arena: str) -> list[list[dict]]:
+    """Facts sharing (subject_entity_id, predicate) but with DIFFERENT
+    statements — candidate same-assertion-different-words pairs for LLM
+    adjudication. Exact-triple dupes already collapse elsewhere; this is the
+    semantic tier. Restricted disclosure excluded.
+
+    RECALL SCOPE (#96 review §5): grouping on (subject, predicate) catches only
+    same-predicate wording variants ("decided X" vs "decided X, confirmed").
+    CROSS-predicate synonyms ("joined Acme" / "works at Acme") are NOT grouped
+    here — that needs a predicate-synonym map or subject-level pairwise
+    adjudication (O(n²), deferred). adjudicate_facts() itself handles any pair;
+    it's the candidate generation that's intentionally narrow to bound LLM calls."""
+    cur.execute(
+        """SELECT id, subject_entity_id, predicate, object_entity_id, statement,
+                  confidence, provenance_event_ids
+           FROM facts WHERE arena=%s AND disclosure_class <> 'restricted'
+             AND subject_entity_id IS NOT NULL""",
+        (arena,))
+    groups: dict[tuple, list[dict]] = {}
+    for f in cur.fetchall():
+        groups.setdefault((f["subject_entity_id"], f["predicate"]), []).append(f)
+    return [g for g in groups.values()
+            if len({_norm(x["statement"]) for x in g}) > 1]  # >1 distinct statement
+
+
+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)")
+    ap.add_argument("--llm-endpoint", default=os.environ.get("PME_V2_LLM_ENDPOINT"),
+                    help="in-VPC distiller /v1/chat/completions for adjudication "
+                         "(no egress). Omit to skip the LLM tier (deterministic only).")
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    args = ap.parse_args()
+    if not args.pg_dsn:
+        print("PG_DSN required", file=sys.stderr)
+        return 2
+    post_fn = _distiller_post_fn(args.llm_endpoint, args.model) if args.llm_endpoint else None
+
+    proposals = merged = llm_entity_merges = llm_fact_merges = 0
+    with psycopg.connect(args.pg_dsn, row_factory=dict_row) as conn:
+        with conn.cursor() as cur:
+            def do_merge(group, signal_note=""):
+                nonlocal proposals, merged
+                sig = _authority_signals(cur, args.arena, [e["id"] for e in group], args.model)
+                master_c, losers_c = C.pick_master(_candidates(group, sig))
+                loser_ids = [l.entity_id for l in losers_c]
+                if not loser_ids:
+                    return 0
+                proposals += 1
+                by_id = {e["id"]: e for e in group}
+                master, losers = by_id[master_c.entity_id], [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{signal_note} → '{master['canonical_name']}' ({master['id']}) "
+                      f"absorbs {[l['canonical_name'] for l in losers]}")
+                if args.apply:
+                    _execute_entity_plan(cur, plan)
+                    _dedup_master_facts(cur, args.arena, master["id"])
+                    conn.commit()
+                merged += len(loser_ids)
+                return len(loser_ids)
+
+            # Tier 1 — deterministic (exact-name + unambiguous cross-run)
+            for group in _entity_dup_sets(cur, args.arena):
+                do_merge(group)
+
+            # Tier 2 — LLM adjudication via the in-VPC distiller (no egress)
+            if post_fn:
+                # 2a. ambiguous cross-run: which real entity does the junk match?
+                for amb in _ambiguous_cross_run(cur, args.arena):
+                    j = amb["junk"]
+                    jctx = _entity_context(cur, args.arena, j["id"])
+                    for cand in amb["candidates"]:
+                        v = adjudicate_entities(
+                            {**j, "context": jctx},
+                            {**cand, "context": _entity_context(cur, args.arena, cand["id"])},
+                            post_fn)
+                        if v["same"]:
+                            print(f"  [llm:{v['reason'][:40]}]", end="")
+                            llm_entity_merges += do_merge([cand, j], signal_note=" (llm)")
+                            break
+                # 2b. semantic fact fusion: same assertion, different words?
+                for fg in _semantic_fact_groups(cur, args.arena):
+                    fg_sorted = sorted(fg, key=lambda f: (f.get("confidence", 0), f["id"]), reverse=True)
+                    keep = fg_sorted[0]
+                    same = [keep]
+                    for other in fg_sorted[1:]:
+                        if adjudicate_facts(keep["statement"], other["statement"], post_fn)["same"]:
+                            same.append(other)
+                    if len(same) > 1:
+                        plan = build_fact_merge_plan(arena=args.arena, dup_facts=same)
+                        print(f"  FACT-MERGE (llm) → '{keep['statement'][:50]}' absorbs {len(same)-1}")
+                        if args.apply and plan:
+                            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,'llm_adjudication',%s,%s::jsonb)""",
+                                    ("fm_" + uuid.uuid4().hex[:20], a["arena"], a["canonical_id"],
+                                     a["deprecated_id"], a["deprecated_statement"], a["provenance_unioned"],
+                                     json.dumps(a["rollback_payload"], default=str)))
+                            cur.execute("DELETE FROM facts WHERE id = ANY(%s)", (plan["deprecated_ids"],))
+                            conn.commit()
+                        llm_fact_merges += len(same) - 1
+
+            run_id = "fdr_" + uuid.uuid4().hex[:20]
+            detail = {"proposals": proposals, "merged": merged,
+                      "llm_entity_merges": llm_entity_merges, "llm_fact_merges": llm_fact_merges,
+                      "llm_tier": bool(post_fn)}
+            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 + llm_fact_merges, json.dumps(detail)))
+        conn.commit()
+
+    label = "APPLY (reversible via entity_merges/fact_merges)" if args.apply else "DRY-RUN"
+    print(f"[fusion-drive:fuse] {label} arena={args.arena}: {proposals} entity proposal(s), "
+          f"{merged} entities merged ({llm_entity_merges} via llm), {llm_fact_merges} facts merged via llm. "
+          f"LLM tier: {'on (distiller)' if post_fn else 'off'}")
+    print(f"  ledger: {run_id}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())