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

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

@@ -37,14 +37,35 @@ 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 required", file=sys.stderr)
+    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())
@@ -101,18 +122,52 @@ def _entity_dup_sets(cur, arena: str) -> list[list[dict]]:
     return exact + cross
 
 
-def _candidates(group: list[dict]) -> list[C.CanonicalCandidate]:
+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 [],
-            # 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.
+            grounded=signals.get(e["id"], {}).get("grounded", False),
+            from_current_teacher=signals.get(e["id"], {}).get("from_current_teacher", False),
         )
         for e in group
     ]
@@ -209,53 +264,168 @@ def _dedup_master_facts(cur, arena: str, master_id: str) -> int:
     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 = 0
-    merged = 0
+    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:
-            groups = _entity_dup_sets(cur, args.arena)
-            for group in groups:
-                master_c, losers_c = C.pick_master(_candidates(group))
+            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:
-                    continue
+                    return 0
                 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]
+                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 → 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)}]")
+                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"])
-                    merged += len(loser_ids)
-                    conn.commit()  # per-merge: a bad merge can't roll back the good ones, and locks stay short
+                    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, json.dumps({"proposals": proposals, "merged": merged})),
-            )
+                 proposals, merged + llm_fact_merges, json.dumps(detail)))
         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")
+    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