@pentatonic-ai/ai-agent-sdk 0.10.4 → 0.10.6

package/packages/memory-engine-v2/resolution-queue-design.md

@@ -0,0 +1,165 @@
+# Design: incremental write-time entity resolution (`resolution_queue`)
+
+**Status:** design only — NO implementation in this PR.
+**Companion:** `scripts/entity_resolution_v2.py` (the one-time/batch
+tool this design makes incremental), `RFC-entity-reconciliation.md`
+(v1: extract-time pairing + alias-aware upsert + backfill).
+
+## Problem
+
+`entity_resolution_v2.py` is batch tooling: an operator runs it
+against one arena, eyeballs the tiered report, and (snapshot-gated)
+applies. But fragmentation is continuous — every day of ingest can
+mint new variant rows ("Johann_Boedecker" from a doc filename,
+"Bödecker, Johann" from a vCard) that v1's exact-surface alias
+resolution at upsert will not catch. Re-running the batch tool
+forever is operationally wrong, and running blocking + embeddings +
+LLM adjudication **on the hot write path is a non-starter**: the
+upsert path is latency-sensitive (advisory-lock held), embeddings are
+a network call, and adjudication is an LLM call.
+
+## Design: mirror `distillation_queue`'s claim pattern
+
+The engine already has exactly the right shape for "expensive
+asynchronous post-processing of a row that was written cheaply":
+`distillation_queue` (001_init.sql) — extractor-sync enqueues, the
+async worker `claim → process → done/failed`, with `claim_expires_at`
+so a crashed worker's items resurface.
+
+`resolution_queue` mirrors it 1:1, keyed on entities instead of
+events:
+
+1. **Enqueue (hot path, cheap):** whenever an extractor inserts a NEW
+   entity row (not an alias-resolved update onto an existing id), it
+   also inserts one `resolution_queue` row inside the same
+   transaction. Cost: one INSERT. No embeddings, no LLM, no blocking
+   on the write path.
+2. **Sweep (nightly, off-path):** a resolver worker claims pending
+   items in batches (same `FOR UPDATE SKIP LOCKED` claim the
+   distiller uses), and for each claimed entity runs the
+   `entity_resolution_v2` pipeline *scoped to that entity*:
+   blocking keys → candidate set within its arena → embedding
+   similarity → threshold bands → LLM adjudication of the ambiguous
+   band → merge via the existing v1 `apply_proposals` machinery
+   (per-proposal transaction, `entity_merges` audit +
+   `rollback_payload`).
+3. **Embeddings cached:** the resolver persists each entity's bundle
+   embedding (and the bundle fingerprint it was computed from) so the
+   nightly sweep only re-embeds entities whose surface forms/facts
+   changed. This is what actually keeps the sweep cheap at 100k+
+   entity scale.
+
+### Policy carried over from the batch tool
+
+- Tier precedence: `co_occurrence > alias_overlap > embedding_llm >
+  heuristic`. The sweep only ever auto-applies the same things the
+  batch tool would auto-apply.
+- `unsure` adjudications NEVER merge — they land in
+  `resolution_queue.status = 'review'` for a human, and stay claimable
+  by a future "review UI" rather than being retried into oblivion.
+- Bare-first-name policy: single-token entities merge only with
+  exactly one block candidate AND adjudication = yes.
+- Arena scoping: the worker processes one claimed row at a time and
+  every statement carries that row's `arena`. A claimed
+  `pentatonic-team` entity can never read or write `pip-agents` rows.
+
+## Migration draft (DESIGN ONLY — not in `org-model/migrations/`)
+
+```sql
+-- DRAFT 004_resolution_queue.sql — do not apply without review.
+
+-- 1. Queue, mirroring distillation_queue's claim pattern.
+CREATE TABLE resolution_queue (
+  id                  BIGSERIAL PRIMARY KEY,
+  entity_id           TEXT NOT NULL REFERENCES entities(id) ON DELETE CASCADE,
+  arena               TEXT NOT NULL,           -- denormalised for scoped claims
+  enqueued_at         TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+  claimed_by          TEXT,
+  claimed_at          TIMESTAMPTZ,
+  claim_expires_at    TIMESTAMPTZ,             -- crashed workers' items resurface
+  status              TEXT NOT NULL DEFAULT 'pending',
+  attempts            INT NOT NULL DEFAULT 0,
+  last_error          TEXT,
+  completed_at        TIMESTAMPTZ,
+  -- 'review' = adjudication returned 'unsure' / bare-name policy hold;
+  -- terminal for the worker, surfaced to a human.
+  CONSTRAINT valid_status CHECK (
+    status IN ('pending', 'claimed', 'done', 'failed', 'review'))
+);
+
+CREATE INDEX idx_resolution_status ON resolution_queue(status);
+CREATE INDEX idx_resolution_arena ON resolution_queue(arena, status);
+CREATE INDEX idx_resolution_entity ON resolution_queue(entity_id);
+CREATE INDEX idx_resolution_claim_expires ON resolution_queue(claim_expires_at)
+  WHERE status = 'claimed';
+
+-- 2. Embedding cache so the sweep doesn't re-embed unchanged entities.
+CREATE TABLE entity_embeddings (
+  entity_id           TEXT PRIMARY KEY REFERENCES entities(id) ON DELETE CASCADE,
+  arena               TEXT NOT NULL,
+  bundle_fingerprint  TEXT NOT NULL,           -- sha256 of the embedded bundle
+  embedding           REAL[] NOT NULL,         -- gateway dim (4096 today)
+  embedding_model     TEXT NOT NULL,
+  embedded_at         TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+CREATE INDEX idx_entity_embeddings_arena ON entity_embeddings(arena);
+
+-- 3. Admit the new merge signal in the audit table. The existing
+--    CHECK only allows ('co_occurrence','alias_overlap','heuristic',
+--    'online_resolver'); entity_resolution_v2.py refuses to apply
+--    embedding_llm-tier proposals until this lands.
+ALTER TABLE entity_merges DROP CONSTRAINT entity_merges_merge_signal_check;
+ALTER TABLE entity_merges ADD CONSTRAINT entity_merges_merge_signal_check
+  CHECK (merge_signal IN (
+    'co_occurrence', 'alias_overlap', 'heuristic',
+    'online_resolver', 'embedding_llm'));
+```
+
+> Note: the CHECK in `002_entity_merges_audit.sql` is unnamed in DDL;
+> Postgres auto-names it `entity_merges_merge_signal_check`. Verify
+> the live name via `pg_constraint` before drafting the final ALTER —
+> the v2 script does this probe programmatically
+> (`_entity_merges_check_allows`).
+
+### Claim query (worker pseudocode)
+
+```sql
+WITH next AS (
+  SELECT id FROM resolution_queue
+  WHERE status = 'pending'
+     OR (status = 'claimed' AND claim_expires_at < NOW())
+  ORDER BY enqueued_at
+  LIMIT 50
+  FOR UPDATE SKIP LOCKED
+)
+UPDATE resolution_queue q SET
+  status = 'claimed', claimed_by = $worker, claimed_at = NOW(),
+  claim_expires_at = NOW() + INTERVAL '15 minutes',
+  attempts = attempts + 1
+FROM next WHERE q.id = next.id
+RETURNING q.id, q.entity_id, q.arena;
+```
+
+## Rollout sequencing
+
+1. Land + run the batch tool (`entity_resolution_v2.py`) once per
+   arena, snapshot-gated — clears the standing backlog. (Per the
+   clean-rebuild runbook: AFTER the distiller-format decision settles,
+   so we don't merge rows the re-distillation would re-shape.)
+2. Apply draft migration 004 (review first; the engine is shared
+   multi-tenant infra — pip-agents stays frozen, the migration adds
+   tables and a CHECK value, touches no rows).
+3. Ship the enqueue (one INSERT in each extractor's entity-insert
+   path) — write-path cost is negligible.
+4. Ship the nightly sweep worker (a sibling of extractor-async,
+   same claim loop), initially in report-only mode (`status='review'`
+   for everything) for a week; then enable auto-apply for the
+   embedding_llm tier once precision on the report holds at ~100%.
+
+## Non-goals
+
+- Real-time resolution at upsert (v1's alias-aware exact-surface
+  resolution already covers the common case there).
+- Cross-arena resolution of any kind. Identity never spans arenas.
+- Non-person entity types (same follow-up posture as the v1 RFC).

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

@@ -53,8 +53,13 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 
-import psycopg
-import psycopg.rows
+try:
+    import psycopg
+    import psycopg.rows
+except ImportError:  # pragma: no cover — allows importing the merge
+    # machinery (entity_resolution_v2.py, unit tests) without the DB
+    # driver installed. main() still requires it to actually run.
+    psycopg = None  # type: ignore[assignment]
 
 
 # ----------------------------------------------------------------------
@@ -533,6 +538,10 @@ def main() -> int:
     if not args.pg_dsn:
         print("error: --pg-dsn (or $PG_DSN) required", file=sys.stderr)
         return 2
+    if psycopg is None:
+        print("error: psycopg is required to run this script "
+              "(pip install 'psycopg[binary]')", file=sys.stderr)
+        return 2
 
     merged_by = args.merged_by or f"backfill-{datetime.now(timezone.utc):%Y-%m}"
 

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

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+"""Backfill: write the named 'lex' BM25 sparse vector onto existing
+points in the Qdrant "evidence" collection (roadmap BET 3, hybrid
+lexical+dense retrieval).
+
+ADDITIVE ONLY. This script writes ONLY the named sparse vector via
+`update_vectors` — the existing unnamed dense vector (Qwen3-Embedding-8B,
+4096-d) is never read, re-embedded, or modified. Qdrant's
+`update_vectors` updates exactly the vector names supplied and leaves
+all other vectors on the point untouched.
+
+Why a backfill at all: Qdrant point payloads only carry a 300-char
+`content_preview` — BM25 term statistics over a truncated preview would
+systematically miss tail terms, so the sparse vector MUST be computed
+from the FULL event content, which lives in Postgres `events.content`
+(joined via the `event_id` payload field).
+
+Pipeline per batch:
+
+  1. scroll(collection, with_payload=["event_id"], with_vectors=False)
+  2. fetch full content: SELECT id, content FROM events WHERE id = ANY(...)
+  3. BM25-encode the full content (fastembed Qdrant/bm25, CPU)
+  4. update_vectors([PointVectors(id=point_id, vector={"lex": sparse})])
+
+Dry-run by default — counts points, resolves content coverage on a
+sample, prints an ETA, writes NOTHING. Pass --apply to write.
+
+Resumable: the last scroll offset (a Qdrant point id) is persisted to
+a state file after every batch; re-running resumes from there.
+Idempotent: re-encoding the same content produces the same sparse
+vector, and update_vectors overwrites in place — re-running over
+already-backfilled points is wasted work, not corruption.
+
+Prerequisites (run when the operator — Phil H — has flipped nothing
+yet, but the collection must already carry the 'lex' sparse vector
+config; either start compat once with SEARCH_HYBRID_ENABLED=1 or run
+this script with --ensure-config):
+
+    python3 backfill_sparse_vectors.py \
+        --qdrant-url http://127.0.0.1:16333 \
+        --pg-dsn postgresql://pme:...@127.0.0.1:15432/org_model \
+        [--apply]                # write; default is dry-run
+        [--ensure-config]        # add the 'lex' sparse config if missing (idempotent)
+        [--arena <arena>]        # optional payload filter
+        [--client <clientId>]    # optional payload filter
+        [--batch-size 256]
+        [--state-file .backfill_sparse_vectors.state.json]
+        [--reset-state]          # ignore + overwrite any previous offset
+        [--max-points N]         # stop after N points (smoke runs)
+
+⚠️ DISK HEADROOM: the sparse index lands on the same Qdrant storage
+volume as the dense vectors. The 2026-06-05 outage was the engine box
+root disk filling up — confirm headroom (sparse 'lex' vectors for the
+~620–745k-point evidence collection are small relative to the 4096-d
+dense data, expect low single-digit GB including index, but CHECK
+`df -h` on the box before --apply).
+
+Exit codes: 0 success; 1 partial failure (some batches errored); 2 bad args.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+
+SPARSE_VECTOR_NAME = "lex"
+SPARSE_MODEL_NAME = os.environ.get("SEARCH_SPARSE_MODEL", "Qdrant/bm25")
+DEFAULT_COLLECTION = "evidence"
+
+# Conservative single-CPU-worker BM25 throughput planning figure for the
+# dry-run ETA (fastembed Qdrant/bm25 tokenize+count is cheap; real rates
+# are usually higher). Override with --eta-rate.
+DEFAULT_ENCODE_RATE_PER_SEC = 400.0
+
+
+# ----------------------------------------------------------------------
+# Pure planning helpers — stdlib only, unit-tested without qdrant/pg.
+# ----------------------------------------------------------------------
+
+
+def batch_count(total_points: int, batch_size: int) -> int:
+    """Number of scroll/encode/update batches for `total_points`."""
+    if total_points <= 0 or batch_size <= 0:
+        return 0
+    return (total_points + batch_size - 1) // batch_size
+
+
+def eta_seconds(total_points: int, rate_per_sec: float) -> float:
+    """Naive ETA: encode dominates (scroll + update are I/O-cheap)."""
+    if total_points <= 0 or rate_per_sec <= 0:
+        return 0.0
+    return total_points / rate_per_sec
+
+
+def format_eta(seconds: float) -> str:
+    seconds = int(seconds)
+    h, rem = divmod(seconds, 3600)
+    m, s = divmod(rem, 60)
+    if h:
+        return f"{h}h{m:02d}m"
+    if m:
+        return f"{m}m{s:02d}s"
+    return f"{s}s"
+
+
+def load_state(path: str) -> dict:
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {}
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def save_state(path: str, state: dict) -> None:
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(state, f, indent=2)
+    os.replace(tmp, path)
+
+
+# ----------------------------------------------------------------------
+# Heavy-dependency sections (lazy imports so the planning helpers above
+# stay importable in dependency-free test environments).
+# ----------------------------------------------------------------------
+
+
+def _build_filter(arena: str | None, client: str | None):
+    from qdrant_client import models as qmodels
+
+    must = []
+    if arena:
+        must.append(qmodels.FieldCondition(key="arena", match=qmodels.MatchValue(value=arena)))
+    if client:
+        must.append(qmodels.FieldCondition(key="clientId", match=qmodels.MatchValue(value=client)))
+    return qmodels.Filter(must=must) if must else None
+
+
+def _ensure_config(qdrant, collection: str) -> bool:
+    """Idempotently add the 'lex' sparse vector config (mirrors compat's
+    _ensure_sparse_vector_config — additive metadata, no point data
+    touched). Returns True if added."""
+    from qdrant_client import models as qmodels
+
+    info = qdrant.get_collection(collection)
+    existing = getattr(info.config.params, "sparse_vectors", None) or {}
+    if SPARSE_VECTOR_NAME in existing:
+        return False
+    qdrant.update_collection(
+        collection_name=collection,
+        sparse_vectors_config={
+            SPARSE_VECTOR_NAME: qmodels.SparseVectorParams(
+                modifier=qmodels.Modifier.IDF,
+                index=qmodels.SparseIndexParams(on_disk=True),
+            )
+        },
+    )
+    return True
+
+
+def run(args: argparse.Namespace) -> int:
+    from qdrant_client import QdrantClient
+    from qdrant_client import models as qmodels
+
+    qdrant = QdrantClient(url=args.qdrant_url, timeout=60)
+    flt = _build_filter(args.arena, args.client)
+
+    # Collection-level facts up front (counts include points outside the
+    # filter; the filtered total is only known after a full scroll).
+    info = qdrant.get_collection(args.collection)
+    total_points = info.points_count or 0
+    sparse_cfg = getattr(info.config.params, "sparse_vectors", None) or {}
+    print(f"[backfill] collection={args.collection} points_count={total_points}")
+    print(f"[backfill] sparse config present: {SPARSE_VECTOR_NAME in sparse_cfg}")
+
+    if SPARSE_VECTOR_NAME not in sparse_cfg:
+        if args.ensure_config:
+            added = _ensure_config(qdrant, args.collection)
+            print(f"[backfill] sparse config '{SPARSE_VECTOR_NAME}' added: {added}")
+        elif args.apply:
+            print(
+                f"error: collection has no '{SPARSE_VECTOR_NAME}' sparse vector config. "
+                "Start compat once with SEARCH_HYBRID_ENABLED=1 or pass --ensure-config.",
+                file=sys.stderr,
+            )
+            return 2
+
+    if not args.apply:
+        eta = eta_seconds(total_points, args.eta_rate)
+        print(
+            f"[backfill] DRY-RUN: would scroll ~{total_points} points "
+            f"in {batch_count(total_points, args.batch_size)} batches of {args.batch_size}; "
+            f"ETA ~{format_eta(eta)} at {args.eta_rate:.0f} pts/s encode rate"
+        )
+
+    # Postgres + encoder only needed past this point.
+    import psycopg
+    import psycopg.rows
+
+    pg = psycopg.connect(args.pg_dsn, row_factory=psycopg.rows.dict_row) if args.pg_dsn else None
+    if pg is None and args.apply:
+        print("error: --pg-dsn (or $PG_DSN) required with --apply", file=sys.stderr)
+        return 2
+
+    encoder = None
+    if args.apply:
+        from fastembed import SparseTextEmbedding
+
+        encoder = SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)
+
+    state = {} if args.reset_state else load_state(args.state_file)
+    offset = state.get("next_offset")
+    if offset:
+        print(f"[backfill] resuming from offset {offset}")
+
+    scanned = 0
+    written = 0
+    skipped_existing = 0
+    missing_content = 0
+    errors = 0
+    t0 = time.monotonic()
+
+    while True:
+        points, next_offset = qdrant.scroll(
+            collection_name=args.collection,
+            scroll_filter=flt,
+            limit=args.batch_size,
+            offset=offset,
+            with_payload=["event_id"],
+            # --apply additionally pulls the (tiny) existing 'lex'
+            # vector so already-backfilled points are skipped on
+            # resume. The dense vector is NEVER fetched.
+            with_vectors=[SPARSE_VECTOR_NAME] if (args.apply and not args.force) else False,
+        )
+        if not points:
+            break
+
+        batch = []  # (point_id, event_id)
+        for p in points:
+            scanned += 1
+            eid = (p.payload or {}).get("event_id")
+            if not eid:
+                missing_content += 1
+                continue
+            if args.apply and not args.force:
+                vecs = p.vector if isinstance(p.vector, dict) else {}
+                if vecs and SPARSE_VECTOR_NAME in vecs:
+                    skipped_existing += 1
+                    continue
+            batch.append((p.id, eid))
+
+        if batch and pg is not None:
+            event_ids = list({eid for _, eid in batch})
+            with pg.cursor() as cur:
+                cur.execute(
+                    "SELECT id, content FROM events WHERE id = ANY(%s)",
+                    (event_ids,),
+                )
+                content_by_id = {r["id"]: r["content"] for r in cur.fetchall()}
+
+            todo = []
+            for pid, eid in batch:
+                content = content_by_id.get(eid)
+                if not content:
+                    missing_content += 1
+                    continue
+                todo.append((pid, content))
+
+            if args.apply and todo:
+                try:
+                    embs = list(encoder.embed([c for _, c in todo]))
+                    qdrant.update_vectors(
+                        collection_name=args.collection,
+                        points=[
+                            qmodels.PointVectors(
+                                id=pid,
+                                vector={
+                                    SPARSE_VECTOR_NAME: qmodels.SparseVector(
+                                        indices=[int(i) for i in e.indices],
+                                        values=[float(v) for v in e.values],
+                                    )
+                                },
+                            )
+                            for (pid, _), e in zip(todo, embs)
+                        ],
+                        wait=True,
+                    )
+                    written += len(todo)
+                except Exception as e:
+                    errors += 1
+                    print(f"[backfill] batch ERROR at offset {offset}: {e}", file=sys.stderr)
+
+        offset = next_offset
+        save_state(args.state_file, {
+            "next_offset": str(offset) if offset is not None else None,
+            "scanned": scanned,
+            "written": written,
+            "mode": "apply" if args.apply else "dry-run",
+            "collection": args.collection,
+        })
+
+        if scanned and scanned % (args.batch_size * 20) == 0:
+            rate = scanned / max(time.monotonic() - t0, 1e-6)
+            remaining = max(total_points - scanned, 0)
+            print(
+                f"[backfill] scanned={scanned} written={written} "
+                f"skipped={skipped_existing} missing_content={missing_content} "
+                f"rate={rate:.0f}/s eta={format_eta(remaining / max(rate, 1e-6))}"
+            )
+
+        if args.max_points and scanned >= args.max_points:
+            print(f"[backfill] --max-points {args.max_points} reached; stopping")
+            break
+        if next_offset is None:
+            break
+
+    dur = time.monotonic() - t0
+    print(
+        f"[backfill] DONE mode={'apply' if args.apply else 'dry-run'} "
+        f"scanned={scanned} written={written} skipped_existing={skipped_existing} "
+        f"missing_content={missing_content} errors={errors} in {format_eta(dur)}"
+    )
+    if not args.apply:
+        print("[backfill] dry-run only; pass --apply to write the 'lex' vectors")
+    if pg is not None:
+        pg.close()
+    return 1 if errors else 0
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument("--qdrant-url", default=os.environ.get("VECTOR_INDEX_URL", "http://127.0.0.1:16333"))
+    p.add_argument("--pg-dsn", default=os.environ.get("PG_DSN"),
+                   help="postgres DSN for full event content; defaults to $PG_DSN")
+    p.add_argument("--collection", default=DEFAULT_COLLECTION)
+    p.add_argument("--apply", action="store_true", help="write; default is dry-run")
+    p.add_argument("--ensure-config", action="store_true",
+                   help="idempotently add the 'lex' sparse vector config if missing")
+    p.add_argument("--arena", default=None, help="optional arena payload filter")
+    p.add_argument("--client", default=None, help="optional clientId payload filter")
+    p.add_argument("--batch-size", type=int, default=256)
+    p.add_argument("--state-file", default=".backfill_sparse_vectors.state.json")
+    p.add_argument("--reset-state", action="store_true")
+    p.add_argument("--force", action="store_true",
+                   help="re-write 'lex' even where it already exists")
+    p.add_argument("--max-points", type=int, default=0,
+                   help="stop after scanning N points (smoke runs)")
+    p.add_argument("--eta-rate", type=float, default=DEFAULT_ENCODE_RATE_PER_SEC,
+                   help="points/sec planning figure for the dry-run ETA")
+    return p.parse_args(argv)
+
+
+def main() -> int:
+    args = parse_args()
+    if args.batch_size <= 0:
+        print("error: --batch-size must be > 0", file=sys.stderr)
+        return 2
+    return run(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())