@pentatonic-ai/ai-agent-sdk 0.10.19 → 0.10.21

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

@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""Prompt-version-drift re-distillation — RFC-decay-and-fusion Part B′.
+
+The third triage verb of the Fusion Drive: regenerate high-value extractions
+that were produced by a *superseded* teacher/prompt by re-running their source
+event through the *current clean* teacher. Fusion converges horizontally, decay
+ages vertically; this refreshes in depth.
+
+TRIGGER (not raw age): an event is stale when `event_distillations` shows it was
+distilled only under a dirty `system_prompt_hash` (pre-clean: bbdaba/f1e0ff/
+ef0647…) and never under the clean one (`6ccfe70f…`, SDK 0.10.19 = #126 + #129).
+
+MECHANISM: re-insert the event into `distillation_queue` (status=pending). The
+existing extractor-async worker + autoscaler do the rest, writing a fresh clean
+extraction AND a clean trace (which `build_retrain_corpus.py` then harvests — one
+loop repairs the graph and grows the retrain corpus).
+
+SUPERSEDENCE (the load-bearing requirement — Part B′4): the store is
+pure-accretion, so a naive re-enqueue makes the clean extraction land *beside*
+the dirty one and WORSENS fragmentation. The durable fix routes supersede
+through Fusion's tombstone machinery (RFC A2/A3), which is not yet live. This
+tool ships the **interim event-scoped supersede** (Part B′4 Open Q): with
+--supersede-facts it deletes only facts whose provenance set is exactly
+{this event} under the old hash — the single-provenance majority, where no
+corroboration is lost — after dumping a rollback payload. Multi-provenance facts
+and ALL entities are left untouched (entity-level fragmentation is fusion/decay's
+job; deleting an entity cascades to relationships and NULLs other facts' FKs).
+
+SAFETY: dry-run by default; --apply gated; --arena REQUIRED (no global runs);
+--limit caps the batch; every destructive op is dumped to a rollback JSONL first;
+a run ledger is written per invocation. Idempotent: skips events that already
+have a pending/claimed queue row or an existing clean-hash distillation.
+
+Usage:
+    # size the work (no writes):
+    python redistill.py --arena 'pentatonic-team%' --dry-run
+    # re-enqueue a bounded slice (non-destructive — accretes until fusion runs):
+    python redistill.py --arena 'pentatonic-team%' --limit 25 --apply
+    # re-enqueue AND interim-supersede single-provenance stale facts:
+    python redistill.py --arena 'pentatonic-team%' --limit 25 --apply --supersede-facts
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from datetime import datetime, timezone
+
+CLEAN_PROMPT_HASH = "6ccfe70f1286a131"  # SDK 0.10.19 (#126 + #129); verify vs worker.SYSTEM_PROMPT_HASH
+
+
+def _connect(dsn: str):
+    import psycopg
+    import psycopg.rows
+    return psycopg.connect(dsn, row_factory=psycopg.rows.dict_row)
+
+
+def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, limit: int,
+           order_by_recency: bool = False):
+    """Stale events: distilled under a dirty prompt, not yet re-distilled clean.
+
+    IDEMPOTENCY SIGNAL = `distillation_traces.system_prompt_hash`, NOT
+    `event_distillations`. The teacher ALWAYS writes a trace (worker.py
+    `_insert_trace`, producer=='teacher'), but `event_distillations` is written
+    ONLY when CASCADE_ENABLED (worker.py line ~2325). In teacher-only mode a
+    re-distilled event gets a clean *trace* but its event_distillations row stays
+    stamped at the old student hash — so keying the clean-check off the ledger
+    would re-select the same events forever. Traces are the reliable per-event
+    prompt-version record here.
+
+    Dirty evidence = a dirty-hash trace OR a dirty-hash event_distillations row
+    (covers student-distilled events, which have a ledger row but no trace)."""
+    params: list = [arena, clean_hash]          # arena, clean (no-clean-trace check)
+    dirty_trace = "AND t2.system_prompt_hash <> %s"
+    params_tail_dirty: list = [clean_hash]
+    if dirty_hashes:
+        dirty_trace = "AND t2.system_prompt_hash = ANY(%s)"
+        params_tail_dirty = [dirty_hashes]
+    # Build param order to match placeholders below.
+    params = [arena, clean_hash] + params_tail_dirty + [clean_hash, limit]
+    # ORDER BY received_at forces a full scan+sort of the whole arena (~270k rows
+    # for pentatonic-team) every call — pathological for bulk/loop work. Default
+    # OFF so the planner early-terminates once LIMIT stale rows are found (most
+    # arena rows ARE stale, so it scans ~LIMIT rows). Recency is only a salience
+    # proxy anyway, moot until the salience column lands (RFC Part D). Opt in with
+    # --order-by-recency for a one-off "freshest first" pass.
+    order_clause = "ORDER BY e.received_at DESC" if order_by_recency else ""
+    cur.execute(
+        f"""
+        WITH stale AS (
+          SELECT e.id AS event_id, e.received_at AS recency
+          FROM events e
+          WHERE e.arena LIKE %s
+            AND NOT EXISTS (                          -- not yet re-distilled clean
+              SELECT 1 FROM distillation_traces t
+              WHERE t.event_id = e.id AND t.system_prompt_hash = %s)
+            AND NOT EXISTS (                          -- not already queued (so loops/batches advance)
+              SELECT 1 FROM distillation_queue q
+              WHERE q.event_id = e.id AND q.status IN ('pending','claimed'))
+            AND (
+              EXISTS (SELECT 1 FROM distillation_traces t2     -- dirty teacher trace
+                      WHERE t2.event_id = e.id {dirty_trace})
+              OR EXISTS (SELECT 1 FROM event_distillations d   -- or dirty student/teacher ledger
+                         WHERE d.event_id = e.id
+                           AND d.system_prompt_hash IS NOT NULL
+                           AND d.system_prompt_hash <> %s)
+            )
+          {order_clause}
+          LIMIT %s
+        )
+        SELECT s.event_id,
+               -- `@> ARRAY[id]` (containment) uses idx_facts_provenance (GIN);
+               -- `id = ANY(col)` does NOT and seq-scans facts per row (2 min+ at
+               -- limit 1000). Must stay @> for the bulk/loop path to be viable.
+               (SELECT count(*) FROM facts f
+                  WHERE f.provenance_event_ids @> ARRAY[s.event_id]) AS fact_n,
+               (SELECT count(*) FROM facts f
+                  WHERE f.provenance_event_ids @> ARRAY[s.event_id]
+                    AND cardinality(f.provenance_event_ids) = 1) AS solo_fact_n,
+               EXISTS (SELECT 1 FROM distillation_queue q
+                       WHERE q.event_id = s.event_id
+                         AND q.status IN ('pending','claimed')) AS in_flight
+        FROM stale s
+        ORDER BY fact_n DESC
+        """,
+        params,
+    )
+    return cur.fetchall()
+
+
+def dump_and_delete_solo_facts(cur, event_id: str, rollback_fh) -> int:
+    """Delete facts whose provenance is exactly {event_id}; dump them first."""
+    solo = "provenance_event_ids @> ARRAY[%s] AND cardinality(provenance_event_ids) = 1"
+    cur.execute(f"SELECT * FROM facts WHERE {solo}", (event_id,))
+    rows = cur.fetchall()
+    for r in rows:
+        rollback_fh.write(json.dumps({"op": "delete_fact", "event_id": event_id,
+                                      "row": _jsonable(r)}, ensure_ascii=False) + "\n")
+    if rows:
+        cur.execute(f"DELETE FROM facts WHERE {solo}", (event_id,))
+    return len(rows)
+
+
+def _jsonable(row: dict) -> dict:
+    out = {}
+    for k, v in row.items():
+        out[k] = v.isoformat() if isinstance(v, datetime) else v
+    return out
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--arena", required=True, help="arena LIKE filter (REQUIRED — scope safety)")
+    ap.add_argument("--clean-hash", default=CLEAN_PROMPT_HASH)
+    ap.add_argument("--dirty-hashes", help="comma-separated hashes to target (default: any != clean)")
+    ap.add_argument("--limit", type=int, default=50, help="max events to re-enqueue (safety cap)")
+    ap.add_argument("--apply", action="store_true", help="actually write (default: dry-run)")
+    ap.add_argument("--supersede-facts", action="store_true",
+                    help="also delete single-provenance stale facts (interim supersede, Part B′4)")
+    ap.add_argument("--order-by-recency", action="store_true",
+                    help="freshest stale events first (forces a full arena scan+sort — slow; off by default)")
+    ap.add_argument("--rollback-dir", default=".", help="dir for rollback + ledger JSONL")
+    ap.add_argument("--pg-dsn", default=os.environ.get("PG_DSN", ""), help="Postgres DSN")
+    args = ap.parse_args()
+
+    if not args.pg_dsn:
+        print("FATAL: no --pg-dsn / PG_DSN", file=sys.stderr)
+        return 2
+    dirty = [h.strip() for h in args.dirty_hashes.split(",")] if args.dirty_hashes else None
+    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+
+    with _connect(args.pg_dsn) as conn:
+        with conn.cursor() as cur:
+            # The org-model postgres container runs with the Docker default 64MB
+            # /dev/shm; a parallel-gather over the facts GIN exhausts the dynamic
+            # shared-memory segment ("DiskFull ... shared memory"). Disable
+            # parallel workers for this session — temp spills go to the (ample)
+            # data disk instead of shm.
+            cur.execute("SET max_parallel_workers_per_gather = 0")
+            cands = triage(cur, args.arena, args.clean_hash, dirty, args.limit,
+                           order_by_recency=args.order_by_recency)
+
+            total = len(cands)
+            in_flight = sum(1 for c in cands if c["in_flight"])
+            actionable = [c for c in cands if not c["in_flight"]]
+            solo = sum(c["solo_fact_n"] for c in actionable)
+            multi = sum(c["fact_n"] - c["solo_fact_n"] for c in actionable)
+            print(json.dumps({
+                "mode": "apply" if args.apply else "dry-run",
+                "arena": args.arena, "clean_hash": args.clean_hash,
+                "candidates": total, "in_flight_skipped": in_flight,
+                "actionable": len(actionable),
+                "stale_facts_solo_provenance": solo,
+                "stale_facts_multi_provenance_LEFT_ALONE": multi,
+                "supersede_facts": args.supersede_facts,
+            }, indent=2))
+
+            if not args.apply:
+                print("\n[dry-run] no writes. Re-run with --apply to re-enqueue"
+                      + (" + supersede solo facts." if args.supersede_facts else "."))
+                return 0
+
+            rb_path = os.path.join(args.rollback_dir, f"redistill_rollback_{stamp}.jsonl")
+            ledger_path = os.path.join(args.rollback_dir, f"redistill_runs_{stamp}.jsonl")
+            enq = deleted = 0
+            with open(rb_path, "w", encoding="utf-8") as rb, \
+                    open(ledger_path, "w", encoding="utf-8") as led:
+                for c in actionable:
+                    eid = c["event_id"]
+                    if args.supersede_facts:
+                        deleted += dump_and_delete_solo_facts(cur, eid, rb)
+                    cur.execute(
+                        "INSERT INTO distillation_queue (event_id, status, attempts) "
+                        "VALUES (%s, 'pending', 0)", (eid,))
+                    enq += 1
+                    led.write(json.dumps({
+                        "event_id": eid, "arena": args.arena,
+                        "old_hash": "dirty", "target_hash": args.clean_hash,
+                        "solo_facts_superseded": c["solo_fact_n"] if args.supersede_facts else 0,
+                        "enqueued_at": stamp,
+                    }, ensure_ascii=False) + "\n")
+                conn.commit()
+
+            print(json.dumps({
+                "applied": True, "re_enqueued": enq, "facts_superseded": deleted,
+                "rollback": rb_path, "ledger": ledger_path,
+            }, indent=2))
+            if deleted:
+                print(f"\nROLLBACK: facts are in {rb_path} — re-INSERT them to undo.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())