@pentatonic-ai/ai-agent-sdk 0.10.29 → 0.10.31

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.cjs CHANGED
@@ -878,7 +878,7 @@ function fireAndForgetEmit(clientConfig, sessionOpts, messages, result, model) {
878
878
  }
879
879
 
880
880
  // src/telemetry.js
881
- var VERSION = "0.10.29";
881
+ var VERSION = "0.10.31";
882
882
  var TELEMETRY_URL = "https://sdk-telemetry.philip-134.workers.dev";
883
883
  function machineId() {
884
884
  const raw = typeof process !== "undefined" ? `${process.env?.USER || process.env?.USERNAME || "u"}:${process.platform || "x"}:${process.arch || "x"}` : "browser";
package/dist/index.js CHANGED
@@ -847,7 +847,7 @@ function fireAndForgetEmit(clientConfig, sessionOpts, messages, result, model) {
847
847
  }
848
848
 
849
849
  // src/telemetry.js
850
- var VERSION = "0.10.29";
850
+ var VERSION = "0.10.31";
851
851
  var TELEMETRY_URL = "https://sdk-telemetry.philip-134.workers.dev";
852
852
  function machineId() {
853
853
  const raw = typeof process !== "undefined" ? `${process.env?.USER || process.env?.USERNAME || "u"}:${process.platform || "x"}:${process.arch || "x"}` : "browser";
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@pentatonic-ai/ai-agent-sdk",
3
- "version": "0.10.29",
3
+ "version": "0.10.31",
4
4
  "description": "TES SDK — LLM observability and lifecycle tracking via Pentatonic Thing Event System. Track token usage, tool calls, and conversations. Manage things through event-sourced lifecycle stages with AI enrichment and vector search.",
5
5
  "type": "module",
6
6
  "main": "./dist/index.cjs",
@@ -131,6 +131,64 @@ STUDENT_MAX_TOKENS = int(os.environ.get("STUDENT_MAX_TOKENS", "768"))
131
131
  # the box's KV-cache headroom. Do NOT raise CONCURRENT_LLM_CALLS to fix student
132
132
  # throughput — that also raises teacher 27B concurrency (OOM risk).
133
133
  STUDENT_CONCURRENT_CALLS = int(os.environ.get("STUDENT_CONCURRENT_CALLS", "48"))
134
+ # How many claimed batches the worker processes CONCURRENTLY. The prior loop
135
+ # claimed a batch, awaited it end-to-end (student gather → straggler tail →
136
+ # teacher tail), and only then claimed the next — a hard barrier. Wall-clock
137
+ # profiling of the 25k-backlog incident (2026-07-08) showed the student L40S
138
+ # burst to Running:50 for ~10s then sit at Running:1 for ~3min draining a
139
+ # straggler tail while the (LB'd) teacher fleet idled — the GPUs were starved by
140
+ # the loop, not by capacity. Overlapping N batches, all sharing ONE student pool
141
+ # and ONE teacher pool (the shared semaphores below), keeps those pools
142
+ # continuously full: batch B's fresh student calls take the slots batch A's
143
+ # stragglers free. Total LLM concurrency is still bounded by the shared sems
144
+ # (STUDENT_CONCURRENT_CALLS / CONCURRENT_LLM_CALLS) REGARDLESS of this value —
145
+ # this only controls how many independent claim→apply pipelines feed them (each
146
+ # on its own PG connection; applies within a batch stay sequential, so one conn
147
+ # per batch is never touched concurrently). Raise to feed the pools harder; the
148
+ # sems, not this, cap GPU load.
149
+ MAX_INFLIGHT_BATCHES = int(os.environ.get("MAX_INFLIGHT_BATCHES", "4"))
150
+ # Process-global shared concurrency pools, created once inside amain's event
151
+ # loop (NOT per-batch — per-batch pools are what caused the burst-then-idle).
152
+ # All in-flight batches acquire from these, so the student/teacher GPUs see one
153
+ # continuous stream of calls instead of per-batch bursts. None until amain runs.
154
+ _STUDENT_SEM: "asyncio.Semaphore | None" = None
155
+ _TEACHER_SEM: "asyncio.Semaphore | None" = None
156
+
157
+ # Cap on provenance_event_ids array length per node (entity/fact/relationship).
158
+ # UNBOUNDED accretion was the drain-throughput bottleneck: hot nodes re-mentioned
159
+ # thousands of times grew arrays to 50k+ event IDs (~2MB TOASTed), so every
160
+ # subsequent upsert read+rewrote the whole array off disk (16s DataFileRead
161
+ # measured; entities TOAST hit 233GB, relationships 66GB — 2026-07-08). Keep only
162
+ # the most-recent N distinct IDs — provenance drives corroboration counts +
163
+ # citations, and the avg node has ~6 (only a handful of super-nodes exceed this),
164
+ # so 200 preserves all real provenance while bounding the pathological tail.
165
+ # Capping the STORED array is self-healing: the first capped write shrinks a
166
+ # bloated node to N, then its upserts stay cheap.
167
+ PROVENANCE_MAX_EVENT_IDS = int(os.environ.get("PROVENANCE_MAX_EVENT_IDS", "200"))
168
+
169
+
170
+ def _cap_prov(existing: str, incoming: str) -> str:
171
+ """SQL expression: merge the `existing` and `incoming` text[] provenance
172
+ arrays, dedupe, and keep the PROVENANCE_MAX_EVENT_IDS most-recent IDs (by
173
+ position — appended IDs are newest), re-ordered oldest→newest to preserve the
174
+ append convention readers see. Drop-in replacement for the old unbounded
175
+ `ARRAY(SELECT DISTINCT UNNEST(existing || incoming))`. `incoming` may carry a
176
+ single %s placeholder; it stays in the same param position."""
177
+ return (
178
+ "(SELECT COALESCE(array_agg(ev ORDER BY ord), ARRAY[]::text[]) "
179
+ "FROM (SELECT ev, max(ord) AS ord "
180
+ f"FROM unnest(({existing}) || ({incoming})) WITH ORDINALITY AS _pu(ev, ord) "
181
+ f"GROUP BY ev ORDER BY max(ord) DESC LIMIT {PROVENANCE_MAX_EVENT_IDS}) _pk)"
182
+ )
183
+
184
+
185
+ # Precomputed capped-provenance SQL fragments for each upsert site (built once).
186
+ _PROV_ENT_UPDATE = _cap_prov("provenance_event_ids", "%s::text[]")
187
+ _PROV_ENT_CONFLICT = _cap_prov("entities.provenance_event_ids", "EXCLUDED.provenance_event_ids")
188
+ _PROV_FACT = _cap_prov("facts.provenance_event_ids", "EXCLUDED.provenance_event_ids")
189
+ _PROV_REL = _cap_prov("relationships.provenance_event_ids", "EXCLUDED.provenance_event_ids")
190
+
191
+
134
192
  # Fraction of student-passing events ALSO sent to the teacher — NOT a quality
135
193
  # lever, this is the monitoring + active-learning sample (drift detection and
136
194
  # fresh teacher-gold on live traffic). Default conservative-ish 5%.
@@ -1888,7 +1946,7 @@ def upsert_entities(
1888
1946
  """
1889
1947
  UPDATE entities SET
1890
1948
  aliases = ARRAY(SELECT DISTINCT UNNEST(aliases || %s::text[])),
1891
- provenance_event_ids = ARRAY(SELECT DISTINCT UNNEST(provenance_event_ids || %s::text[])),
1949
+ provenance_event_ids = """ + _PROV_ENT_UPDATE + """,
1892
1950
  -- Merge the org-domain hard key (no-op for the `{}` default
1893
1951
  -- on non-orgs / unmatched orgs; never clobbers an existing key).
1894
1952
  attributes = entities.attributes || %s::jsonb,
@@ -1964,9 +2022,7 @@ def upsert_entities(
1964
2022
  aliases = (
1965
2023
  SELECT ARRAY(SELECT DISTINCT UNNEST(entities.aliases || EXCLUDED.aliases))
1966
2024
  ),
1967
- provenance_event_ids = (
1968
- SELECT ARRAY(SELECT DISTINCT UNNEST(entities.provenance_event_ids || EXCLUDED.provenance_event_ids))
1969
- ),
2025
+ provenance_event_ids = """ + _PROV_ENT_CONFLICT + """,
1970
2026
  -- merge the org-domain hard key (no-op for `{}`)
1971
2027
  attributes = entities.attributes || EXCLUDED.attributes,
1972
2028
  -- re-corroboration can only RAISE salience, never lower it
@@ -2085,11 +2141,7 @@ def upsert_facts(
2085
2141
  COALESCE(%s, NOW()), %s, %s
2086
2142
  )
2087
2143
  ON CONFLICT (id) DO UPDATE SET
2088
- provenance_event_ids = (
2089
- SELECT ARRAY(SELECT DISTINCT UNNEST(
2090
- facts.provenance_event_ids || EXCLUDED.provenance_event_ids
2091
- ))
2092
- ),
2144
+ provenance_event_ids = """ + _PROV_FACT + """,
2093
2145
  salience = GREATEST(facts.salience, EXCLUDED.salience),
2094
2146
  -- Confidence bumps with each additional independent
2095
2147
  -- source. The cardinality of the merged provenance
@@ -2200,9 +2252,7 @@ def upsert_relationships(
2200
2252
  -- distillation + the SEE-490 backfill on merged-entity edges.
2201
2253
  ON CONFLICT (arena, from_entity_id, to_entity_id, relationship_type) DO UPDATE SET
2202
2254
  weight = relationships.weight + EXCLUDED.weight,
2203
- provenance_event_ids = (
2204
- SELECT ARRAY(SELECT DISTINCT UNNEST(relationships.provenance_event_ids || EXCLUDED.provenance_event_ids))
2205
- ),
2255
+ provenance_event_ids = """ + _PROV_REL + """,
2206
2256
  attributes = relationships.attributes || EXCLUDED.attributes,
2207
2257
  last_seen = GREATEST(relationships.last_seen, EXCLUDED.last_seen),
2208
2258
  first_seen = LEAST(relationships.first_seen, EXCLUDED.first_seen)
@@ -2662,7 +2712,11 @@ async def _process_cascade(
2662
2712
  # STUDENT_CONCURRENT_CALLS, not CONCURRENT_LLM_CALLS: the student is a cheap
2663
2713
  # single-event 4B call that batches far higher than the teacher's 27B
2664
2714
  # batch-calls (see the config note). Reusing the teacher knob starved it.
2665
- sem = asyncio.Semaphore(STUDENT_CONCURRENT_CALLS)
2715
+ # SHARED across all in-flight batches (see MAX_INFLIGHT_BATCHES) so the
2716
+ # student pool stays continuously full — a per-batch semaphore burst-then-
2717
+ # idled the GPU. amain initialises it before any batch runs.
2718
+ sem = _STUDENT_SEM
2719
+ assert sem is not None, "amain must initialise _STUDENT_SEM before processing"
2666
2720
 
2667
2721
  async def _student(item):
2668
2722
  async with sem:
@@ -2719,8 +2773,12 @@ async def _run_teacher(
2719
2773
  chunks.append((chunk_items, chunk_events))
2720
2774
 
2721
2775
  # Bound parallelism with a semaphore — gather only waits for slots,
2722
- # not for everything to be ready at once.
2723
- sem = asyncio.Semaphore(CONCURRENT_LLM_CALLS)
2776
+ # not for everything to be ready at once. SHARED across all in-flight
2777
+ # batches so the teacher fleet (behind the LB) sees a continuous stream and
2778
+ # its total concurrency stays capped at CONCURRENT_LLM_CALLS regardless of
2779
+ # how many batches are escalating at once. amain initialises it.
2780
+ sem = _TEACHER_SEM
2781
+ assert sem is not None, "amain must initialise _TEACHER_SEM before processing"
2724
2782
 
2725
2783
  async def run_one(chunk_items, chunk_events):
2726
2784
  async with sem:
@@ -2917,20 +2975,65 @@ async def amain():
2917
2975
  if stub_mode:
2918
2976
  log.warning("LLM_ENDPOINT not set — running in stub mode (no extraction).")
2919
2977
 
2978
+ # Shared concurrency pools — created HERE, inside the running event loop, so
2979
+ # ONE student pool and ONE teacher pool are shared by every in-flight batch
2980
+ # (see MAX_INFLIGHT_BATCHES). Per-batch pools burst the GPU to Running:~50
2981
+ # then idled it at Running:1 through each batch's straggler/teacher tail; one
2982
+ # shared pool stays continuously full across overlapping batches.
2983
+ global _STUDENT_SEM, _TEACHER_SEM
2984
+ _STUDENT_SEM = asyncio.Semaphore(STUDENT_CONCURRENT_CALLS)
2985
+ _TEACHER_SEM = asyncio.Semaphore(CONCURRENT_LLM_CALLS)
2986
+ log.info(
2987
+ f"pipelined loop: max_inflight_batches={MAX_INFLIGHT_BATCHES}, "
2988
+ f"student_pool={STUDENT_CONCURRENT_CALLS}, teacher_pool={CONCURRENT_LLM_CALLS}"
2989
+ )
2990
+
2920
2991
  # Single async client across the loop. The httpx default limits
2921
- # (max_connections=100, max_keepalive=20) easily cover BATCH_SIZE
2992
+ # (max_connections=100, max_keepalive=20) easily cover the shared-pool
2922
2993
  # concurrent inflight LLM calls.
2923
2994
  async with httpx.AsyncClient(timeout=LLM_TIMEOUT_SEC) as http:
2995
+
2996
+ async def _run_one_batch() -> int:
2997
+ """Claim + fully process one batch on a DEDICATED connection.
2998
+ Returns the number of items claimed (0 = queue empty). Applies within
2999
+ a batch happen sequentially (after each gather), so a batch's
3000
+ connection is never touched concurrently — one conn per batch is
3001
+ safe. The `with` closes the conn on both success and exception; a
3002
+ crash leaves that batch's items claimed until CLAIM_TTL_SEC expires
3003
+ and they are reclaimed (unchanged from the prior loop)."""
3004
+ with psycopg.connect(PG_DSN, autocommit=True) as conn:
3005
+ items = claim_next_batch(conn)
3006
+ if not items:
3007
+ return 0
3008
+ log.info(f"claimed {len(items)} item(s)")
3009
+ await process_batch(http, conn, items, stub_mode)
3010
+ return len(items)
3011
+
3012
+ # Keep up to MAX_INFLIGHT_BATCHES independent claim→apply pipelines
3013
+ # running at once. Their student/teacher calls all draw from the shared
3014
+ # pools above, so while one batch drains its slow tail the others keep
3015
+ # the GPUs fed. Total GPU concurrency is capped by the sems, NOT by this
3016
+ # count.
3017
+ inflight: set[asyncio.Task] = set()
2924
3018
  while True:
2925
3019
  try:
2926
- with psycopg.connect(PG_DSN, autocommit=True) as conn:
2927
- items = claim_next_batch(conn)
2928
- if not items:
2929
- await asyncio.sleep(POLL_INTERVAL_SEC)
2930
- continue
2931
-
2932
- log.info(f"claimed {len(items)} item(s)")
2933
- await process_batch(http, conn, items, stub_mode)
3020
+ while len(inflight) < MAX_INFLIGHT_BATCHES:
3021
+ inflight.add(asyncio.create_task(_run_one_batch()))
3022
+ done, inflight = await asyncio.wait(
3023
+ inflight, return_when=asyncio.FIRST_COMPLETED
3024
+ )
3025
+ claimed_any = False
3026
+ for t in done:
3027
+ try:
3028
+ if t.result() > 0:
3029
+ claimed_any = True
3030
+ except Exception as exc:
3031
+ log.error(f"batch task error: {exc}")
3032
+ # All finished batches were empty claims ⇒ queue drained; back
3033
+ # off before re-claiming so idle workers don't hot-spin. Under a
3034
+ # real backlog every claim returns items and this never sleeps.
3035
+ if not claimed_any:
3036
+ await asyncio.sleep(POLL_INTERVAL_SEC)
2934
3037
  except Exception as exc:
2935
3038
  log.error(f"worker loop error: {exc}")
2936
3039
  await asyncio.sleep(POLL_INTERVAL_SEC * 2)
@@ -49,6 +49,30 @@ log = logging.getLogger("extractor-sync")
49
49
 
50
50
  PG_DSN = os.environ.get("PG_DSN", "postgresql://pme:local-dev-pw@org-model:5432/org_model")
51
51
 
52
+ # Cap on provenance_event_ids array length. Unbounded accretion bloated the
53
+ # entities table to 249GB (233GB TOAST) — hot nodes grew 50k+-element arrays so
54
+ # every upsert read+rewrote ~2MB off disk (16s DataFileRead, 2026-07-08). Keep
55
+ # only the most-recent N distinct IDs. Byte-identical intent to extractor-async's
56
+ # _cap_prov (kept as a local copy per the sibling-service duplication convention,
57
+ # like confidence.py / source_time.py).
58
+ PROVENANCE_MAX_EVENT_IDS = int(os.environ.get("PROVENANCE_MAX_EVENT_IDS", "200"))
59
+
60
+
61
+ def _cap_prov(existing: str, incoming: str) -> str:
62
+ """Merge `existing`||`incoming` text[] provenance, dedupe, keep the
63
+ PROVENANCE_MAX_EVENT_IDS most-recent IDs (append order = newest last).
64
+ Drop-in for the old unbounded `ARRAY(SELECT DISTINCT UNNEST(...))`."""
65
+ return (
66
+ "(SELECT COALESCE(array_agg(ev ORDER BY ord), ARRAY[]::text[]) "
67
+ "FROM (SELECT ev, max(ord) AS ord "
68
+ f"FROM unnest(({existing}) || ({incoming})) WITH ORDINALITY AS _pu(ev, ord) "
69
+ f"GROUP BY ev ORDER BY max(ord) DESC LIMIT {PROVENANCE_MAX_EVENT_IDS}) _pk)"
70
+ )
71
+
72
+
73
+ _PROV_ENT_UPDATE = _cap_prov("provenance_event_ids", "%s::text[]")
74
+ _PROV_ENT_CONFLICT = _cap_prov("entities.provenance_event_ids", "EXCLUDED.provenance_event_ids")
75
+
52
76
  # Connection pool. Sized for high-concurrency bursts: with TES at
53
77
  # shardCount=8 and BATCH_SIZE=50 in the DO drain, compat's parallel
54
78
  # asyncio.gather in /store-batch can fan ~400 simultaneous _extract
@@ -534,7 +558,7 @@ async def _upsert_entities(
534
558
  """
535
559
  UPDATE entities SET
536
560
  aliases = ARRAY(SELECT DISTINCT UNNEST(aliases || %s::text[])),
537
- provenance_event_ids = ARRAY(SELECT DISTINCT UNNEST(provenance_event_ids || %s::text[])),
561
+ provenance_event_ids = """ + _PROV_ENT_UPDATE + """,
538
562
  -- Widen the seen-window with this event's SOURCE time,
539
563
  -- not NOW(): newest evidence = newest source time.
540
564
  last_seen = GREATEST(last_seen, COALESCE(%s, NOW())),
@@ -574,9 +598,7 @@ async def _upsert_entities(
574
598
  aliases = (
575
599
  SELECT ARRAY(SELECT DISTINCT UNNEST(entities.aliases || EXCLUDED.aliases))
576
600
  ),
577
- provenance_event_ids = (
578
- SELECT ARRAY(SELECT DISTINCT UNNEST(entities.provenance_event_ids || EXCLUDED.provenance_event_ids))
579
- ),
601
+ provenance_event_ids = """ + _PROV_ENT_CONFLICT + """,
580
602
  -- merge the email/domain hard key (no-op for the `{}` default
581
603
  -- on entity types that don't stamp it); never clobbers.
582
604
  attributes = entities.attributes || EXCLUDED.attributes,
@@ -17,6 +17,13 @@ from __future__ import annotations
17
17
  from dataclasses import dataclass, field
18
18
  from typing import Any
19
19
 
20
+ # Defensive bound on merged provenance arrays. A merge of many fragmented
21
+ # entities could otherwise union thousands of provenance IDs into the master,
22
+ # re-bloating the entities/relationships TOAST that the extractor cap fixes
23
+ # (see extractor-async/worker.py PROVENANCE_MAX_EVENT_IDS). Kept a literal here
24
+ # to preserve this module's no-env/no-clock purity; must match the extractor.
25
+ PROVENANCE_MAX_EVENT_IDS = 200
26
+
20
27
 
21
28
  # ── entity fusion plan ───────────────────────────────────────────────
22
29
  @dataclass
@@ -75,7 +82,7 @@ def build_entity_merge_plan(
75
82
  provenance = _union(
76
83
  master.get("provenance_event_ids", []),
77
84
  *[l.get("provenance_event_ids", []) for l in losers],
78
- )
85
+ )[:PROVENANCE_MAX_EVENT_IDS]
79
86
 
80
87
  plan = EntityMergePlan(
81
88
  arena=arena,
@@ -112,7 +119,7 @@ def build_entity_merge_plan(
112
119
  "drop": r["id"],
113
120
  "summed_weight": round(keep.get("weight", 1.0) + r.get("weight", 1.0), 4),
114
121
  "provenance": _union(keep.get("provenance_event_ids", []),
115
- r.get("provenance_event_ids", [])),
122
+ r.get("provenance_event_ids", []))[:PROVENANCE_MAX_EVENT_IDS],
116
123
  })
117
124
  else:
118
125
  by_key[key] = r
@@ -147,7 +154,7 @@ def build_fact_merge_plan(*, arena: str, dup_facts: list[dict]) -> dict | None:
147
154
  )
148
155
  master, losers = ranked[0], ranked[1:]
149
156
  provenance = _union(master.get("provenance_event_ids", []),
150
- *[l.get("provenance_event_ids", []) for l in losers])
157
+ *[l.get("provenance_event_ids", []) for l in losers])[:PROVENANCE_MAX_EVENT_IDS]
151
158
  return {
152
159
  "arena": arena,
153
160
  "master_id": master["id"],
@@ -0,0 +1,115 @@
1
+ #!/usr/bin/env python3
2
+ """cap_provenance.py — one-time truncation of over-long provenance_event_ids
3
+ arrays on the graph PG (org_model).
4
+
5
+ Unbounded provenance accretion bloated `entities` to 249GB (233GB TOAST) and
6
+ `relationships` to 72GB (66GB TOAST): a handful of hot nodes grew arrays to
7
+ 50k+ event IDs (~2MB TOASTed), so every upsert read+rewrote the whole array off
8
+ disk — a 16s DataFileRead per update, which was THE distillation drain-throughput
9
+ bottleneck (2026-07-08). The extractor cap (worker.py / server.py
10
+ PROVENANCE_MAX_EVENT_IDS) stops NEW growth; this script shrinks the arrays that
11
+ are ALREADY over the cap so their next upsert is cheap immediately, rather than
12
+ waiting for each hot node to be re-mentioned.
13
+
14
+ It keeps the most-recent CAP distinct IDs (same recency-preserving expression as
15
+ the extractor write path). Runs in bounded batches with a commit per batch so it
16
+ never holds a long lock against live distillation, and is fully idempotent
17
+ (re-running skips already-capped rows via the cardinality guard).
18
+
19
+ DRY-RUN by default (reports how many rows would be capped, writes nothing).
20
+ Pass --apply to write.
21
+
22
+ # in the extractor-async container (has psycopg + PG_DSN):
23
+ python3 cap_provenance.py --table entities --apply
24
+ python3 cap_provenance.py --table relationships --apply
25
+
26
+ Disk reclaim: truncation turns the old 2MB array versions into dead TOAST
27
+ tuples. autovacuum frees them for REUSE (so the table stops growing) but does
28
+ NOT shrink the file. To return the ~299GB to disk, run VACUUM FULL (locks) or
29
+ pg_repack (online) in a maintenance window — non-urgent; performance recovers
30
+ as soon as the live arrays are small.
31
+ """
32
+ from __future__ import annotations
33
+
34
+ import argparse
35
+ import os
36
+ import sys
37
+ import time
38
+
39
+ import psycopg
40
+
41
+ # Same recency-preserving cap as the extractor write path (worker.py._cap_prov):
42
+ # dedupe existing, keep the CAP most-recent by position, re-order oldest→newest.
43
+ _CAP_SQL = (
44
+ "(SELECT COALESCE(array_agg(ev ORDER BY ord), ARRAY[]::text[]) "
45
+ "FROM (SELECT ev, max(ord) AS ord "
46
+ "FROM unnest(provenance_event_ids) WITH ORDINALITY AS _pu(ev, ord) "
47
+ "GROUP BY ev ORDER BY max(ord) DESC LIMIT %(cap)s) _pk)"
48
+ )
49
+
50
+ ALLOWED_TABLES = {"entities", "facts", "relationships"}
51
+
52
+
53
+ def main() -> int:
54
+ ap = argparse.ArgumentParser()
55
+ ap.add_argument("--table", required=True, choices=sorted(ALLOWED_TABLES))
56
+ ap.add_argument("--cap", type=int, default=int(os.environ.get("PROVENANCE_MAX_EVENT_IDS", "200")))
57
+ ap.add_argument("--batch", type=int, default=200, help="ids updated per commit")
58
+ ap.add_argument("--apply", action="store_true", help="write (default: dry-run)")
59
+ args = ap.parse_args()
60
+
61
+ # choices= already guards this; belt-and-braces since the table name is
62
+ # interpolated into DDL-ish SQL (can't be a bound parameter).
63
+ if args.table not in ALLOWED_TABLES:
64
+ print(f"refusing unknown table {args.table!r}", file=sys.stderr)
65
+ return 2
66
+ tbl = args.table
67
+
68
+ dsn = os.environ["PG_DSN"]
69
+ with psycopg.connect(dsn, autocommit=True) as conn:
70
+ over = conn.execute(
71
+ f"SELECT count(*) FROM {tbl} WHERE cardinality(provenance_event_ids) > %s",
72
+ (args.cap,),
73
+ ).fetchone()[0]
74
+ mx = conn.execute(
75
+ f"SELECT COALESCE(max(cardinality(provenance_event_ids)), 0) FROM {tbl}"
76
+ ).fetchone()[0]
77
+ print(f"{tbl}: {over} row(s) over cap={args.cap} (current max cardinality={mx})")
78
+ if not args.apply:
79
+ print("dry-run — pass --apply to truncate")
80
+ return 0
81
+ if over == 0:
82
+ print("nothing to do")
83
+ return 0
84
+
85
+ ids = [
86
+ r[0]
87
+ for r in conn.execute(
88
+ f"SELECT id FROM {tbl} WHERE cardinality(provenance_event_ids) > %s",
89
+ (args.cap,),
90
+ ).fetchall()
91
+ ]
92
+ total = len(ids)
93
+ done = 0
94
+ t0 = time.monotonic()
95
+ for i in range(0, total, args.batch):
96
+ chunk = ids[i : i + args.batch]
97
+ # Re-check cardinality in the WHERE so a concurrent write that already
98
+ # shrank a row is skipped (idempotent, no wasted rewrite).
99
+ conn.execute(
100
+ f"UPDATE {tbl} SET provenance_event_ids = {_CAP_SQL} "
101
+ f"WHERE id = ANY(%(ids)s) AND cardinality(provenance_event_ids) > %(cap)s",
102
+ {"ids": chunk, "cap": args.cap},
103
+ )
104
+ done += len(chunk)
105
+ print(f" capped {done}/{total} ({time.monotonic()-t0:.0f}s)", flush=True)
106
+
107
+ mx2 = conn.execute(
108
+ f"SELECT COALESCE(max(cardinality(provenance_event_ids)), 0) FROM {tbl}"
109
+ ).fetchone()[0]
110
+ print(f"done: {total} row(s) capped; max cardinality now {mx2} (<= {args.cap})")
111
+ return 0
112
+
113
+
114
+ if __name__ == "__main__":
115
+ raise SystemExit(main())