pi-mega-compact 0.7.6 → 0.7.8

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.
@@ -24,6 +24,10 @@ import {
24
24
  latestModelSnapshot,
25
25
  upsertRepoRegistry,
26
26
  recordRepoModel,
27
+ getDedupStats,
28
+ getCompactCount,
29
+ getRecallInjected,
30
+ getCacheHitTokensSaved,
27
31
  type ModelSnapshot,
28
32
  } from "../src/store/sqlite.js";
29
33
  import { detectCrossRepoDrift } from "../src/driftDetection.js";
@@ -73,6 +77,11 @@ interface SessionRuntime {
73
77
  tokensSaved: number; // this session-instance only: reset on session_start
74
78
  lastCompactAt: number | null; // wall-clock ms of the last compaction this session
75
79
  lastNativeCompactAt: number | null; // COMPACT-DEDUP FIX: wall-clock ms of the last NATIVE pi compaction (session_compact event) — used by the agent_end/legacy race guard to skip a redundant ctx.compact() that would throw "Already compacted".
80
+ // S25: live dashboard counters (reset on session_start, mirrored to SQLite).
81
+ compactCount: number; // compactions performed this session-instance
82
+ recallInjections: number; // recall blocks injected this session-instance
83
+ cacheHitTokens: number; // tokens saved via cache hits (dedup + recall) this session
84
+ lengthStopPending: boolean; // S28: set on turn_end when stopReason==='length'
76
85
  }
77
86
 
78
87
  /** ANSI palette for the toolbar. The pi TUI's Text component preserves ANSI
@@ -94,6 +103,11 @@ export const C = {
94
103
 
95
104
  const PULSE = ["◐", "◓", "◑", "◒"];
96
105
 
106
+ // Rough tokens-processed-per-second heuristic for the dashboard's "time saved"
107
+ // estimate. Throughput varies by model/hardware; this is order-of-magnitude so
108
+ // the dashboard can show a human-readable figure, not a precise measurement.
109
+ const TOKENS_PER_SEC_ESTIMATE = 2000;
110
+
97
111
  // ── Full-width widget panel helpers ────────────────────────────────────────
98
112
  // pi's above-editor widget renderer (a Container of Text lines) does NOT pass
99
113
  // a terminal width to setWidget(), so lines render left-aligned by default. To
@@ -274,6 +288,10 @@ export class MegaRuntime {
274
288
  tokensSaved: 0,
275
289
  lastCompactAt: null,
276
290
  lastNativeCompactAt: null,
291
+ compactCount: 0,
292
+ recallInjections: 0,
293
+ cacheHitTokens: 0,
294
+ lengthStopPending: false,
277
295
  };
278
296
  debounceUntil = 0;
279
297
  // S16: debounce for the agent_end resume nudge (avoid busy-loops).
@@ -485,6 +503,12 @@ export class MegaRuntime {
485
503
  const st = this.store.stats(this.rt.sessionId);
486
504
  const repo = this.store.repoStats();
487
505
  const di = this.store.dataInvariant();
506
+ // Live + store-wide cache-hit / compaction counters for the dashboard.
507
+ const ds = getDedupStats(this.currentStateDir);
508
+ const cacheHitsTotal = ds.deduped + getRecallInjected(this.currentStateDir);
509
+ const cacheHitsTotalTokens = getCacheHitTokensSaved(this.currentStateDir);
510
+ const cacheHitsSession = this.rt.dedupSkips + this.rt.recallInjections;
511
+ const sec = (tok: number) => (tok || 0) / TOKENS_PER_SEC_ESTIMATE;
488
512
  // Active model/provider for the current-repo card + the multi-repo table.
489
513
  const modelSnap = latestModelSnapshot(this.currentStateDir);
490
514
  const model = modelSnap
@@ -497,9 +521,13 @@ export class MegaRuntime {
497
521
  }
498
522
  : undefined;
499
523
  // effectiveThresholdPct: the live fire point as a % of the window (null for
500
- // `custom`, which has no tierPct). Used by armed/ready + the dashboard.
524
+ // `custom`, which has no tierPct). S29: honors MEGACOMPACT_AUTO_PCT_TRIGGER
525
+ // override so the dashboard's armed/ready match the context-handler gate
526
+ // (which fires on this same %). Used by armed/ready + the dashboard.
501
527
  const effectiveThresholdPct =
502
- this.config.tierPct != null ? this.config.tierPct * 100 : null;
528
+ this.config.tierPct != null
529
+ ? (this.config.autoPctTrigger ?? this.config.tierPct) * 100
530
+ : null;
503
531
  // armed lights at/above the REAL fire point: max(effectiveThresholdPct,
504
532
  // fastGatePct). fastGatePct already equals tierPct*100 by default, but a
505
533
  // MEGACOMPACT_FAST_GATE_PCT override can raise it, so we take the max.
@@ -507,7 +535,15 @@ export class MegaRuntime {
507
535
  this.lastCtxPercent != null &&
508
536
  this.lastCtxPercent >=
509
537
  Math.max(effectiveThresholdPct ?? 0, this.config.fastGatePct);
510
- const ready = armed && (this.lastCtxTokens ?? 0) >= this.effectiveThreshold;
538
+ // S29: ready mirrors the context-handler gate's basis percent for tiered
539
+ // (the gate fires on pct), tokens for custom (the gate fires on tokens).
540
+ // Previously this always required tokens, so the dashboard could show
541
+ // "armed" (percent high) but never "ready" when tokens were under-reported
542
+ // — the same inconsistency the S29 gate fix removes.
543
+ const ready =
544
+ this.config.tierPct != null
545
+ ? armed && (this.lastCtxPercent ?? 0) >= (effectiveThresholdPct ?? 0)
546
+ : armed && (this.lastCtxTokens ?? 0) >= this.effectiveThreshold;
511
547
  this.dashboard.snapshot({
512
548
  version: 1,
513
549
  updatedAt: new Date().toISOString(),
@@ -604,6 +640,20 @@ export class MegaRuntime {
604
640
  duplicatesCollapsed: di.duplicatesCollapsed,
605
641
  bytesPermanentlyDeleted: di.bytesPermanentlyDeleted,
606
642
  },
643
+ cacheHits: {
644
+ session: cacheHitsSession,
645
+ total: cacheHitsTotal,
646
+ sessionTokensSaved: this.rt.cacheHitTokens,
647
+ totalTokensSaved: cacheHitsTotalTokens,
648
+ },
649
+ compacts: {
650
+ session: this.rt.compactCount,
651
+ total: getCompactCount(this.currentStateDir),
652
+ },
653
+ timeSaved: {
654
+ compact: { sessionSec: sec(this.rt.tokensSaved), totalSec: sec(this.store.repoStats().tokensSaved) },
655
+ cacheHit: { sessionSec: sec(this.rt.cacheHitTokens), totalSec: sec(cacheHitsTotalTokens) },
656
+ },
607
657
  model,
608
658
  } as DashboardSnapshot);
609
659
 
@@ -620,7 +670,9 @@ export class MegaRuntime {
620
670
  : "?";
621
671
  const pctStr =
622
672
  this.lastCtxPercent != null
623
- ? `${Math.round(this.lastCtxPercent * 10) / 10}%`
673
+ ? this.lastCtxPercent > 100
674
+ ? `>100%` // S29: overshoot warning, not a raw "250%" — the percent trigger now compacts before 100%, so this is the residual case where it can't keep up.
675
+ : `${Math.round(this.lastCtxPercent * 10) / 10}%`
624
676
  : "?%";
625
677
  // S24: the tier label is the LIVE pressure band (low/medium/high/ultra/
626
678
  // mega), not the static env preset. It climbs as context fills.
@@ -849,7 +901,11 @@ export class MegaRuntime {
849
901
  tokensSaved: 0,
850
902
  lastCompactAt: null,
851
903
  lastNativeCompactAt: null,
852
- };
904
+ compactCount: 0,
905
+ recallInjections: 0,
906
+ cacheHitTokens: 0,
907
+ lengthStopPending: false,
908
+ };
853
909
  this.statusKey = undefined;
854
910
  this.activeAgents = 0;
855
911
  this.currentTurn = 0;
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "pi-mega-compact",
3
- "version": "0.7.6",
3
+ "version": "0.7.8",
4
4
  "description": "Layered, local, vector-backed context compressor for pi — supersede/collapse/cluster compaction with deduped inline recall.",
5
5
  "type": "module",
6
6
  "license": "BSD-2-Clause",
@@ -0,0 +1,68 @@
1
+ /**
2
+ * sqlite.cachehit.test.ts — tests for the live dashboard counters
3
+ * (incCompactCount / getCompactCount, incRecallInjected / getRecallInjected,
4
+ * incCacheHitTokens / getCacheHitTokensSaved). These reuse the schemaless
5
+ * `meta` integer counter, so no migration is required and the same on-disk
6
+ * SQLite store persists the tallies across (re)opens.
7
+ */
8
+ import { describe, it, beforeEach, afterEach } from "node:test";
9
+ import assert from "node:assert/strict";
10
+ import { mkdtempSync, rmSync } from "node:fs";
11
+ import { tmpdir } from "node:os";
12
+ import { join } from "node:path";
13
+ import {
14
+ openStore,
15
+ incCompactCount,
16
+ getCompactCount,
17
+ incRecallInjected,
18
+ getRecallInjected,
19
+ incCacheHitTokens,
20
+ getCacheHitTokensSaved,
21
+ } from "./sqlite.js";
22
+
23
+ describe("live dashboard counters (meta integers)", () => {
24
+ let dir: string;
25
+ beforeEach(() => {
26
+ dir = mkdtempSync(join(tmpdir(), "cachehit-test-"));
27
+ });
28
+ afterEach(() => {
29
+ rmSync(dir, { recursive: true, force: true });
30
+ });
31
+
32
+ it("increments + reads compact_count and persists on reopen", () => {
33
+ incCompactCount(dir);
34
+ incCompactCount(dir);
35
+ assert.equal(getCompactCount(dir), 2);
36
+ // Ensure the store handle is (re)opened and the value is read back from disk.
37
+ openStore(dir);
38
+ assert.equal(getCompactCount(dir), 2);
39
+ });
40
+
41
+ it("accumulates recall injections + cache-hit tokens", () => {
42
+ incRecallInjected(3, dir);
43
+ incRecallInjected(2, dir);
44
+ assert.equal(getRecallInjected(dir), 5);
45
+ incCacheHitTokens(1200, dir);
46
+ incCacheHitTokens(800, dir);
47
+ assert.equal(getCacheHitTokensSaved(dir), 2000);
48
+ });
49
+
50
+ it("ignores non-positive increments (no-op)", () => {
51
+ incRecallInjected(0, dir);
52
+ incRecallInjected(-5, dir);
53
+ incCacheHitTokens(0, dir);
54
+ assert.equal(getRecallInjected(dir), 0);
55
+ assert.equal(getCacheHitTokensSaved(dir), 0);
56
+ });
57
+
58
+ it("persists all three counters across a fresh openStore handle", () => {
59
+ incCompactCount(dir);
60
+ incRecallInjected(4, dir);
61
+ incCacheHitTokens(500, dir);
62
+ // openStore returns the canonical (on-disk) handle; reading back proves durability.
63
+ openStore(dir);
64
+ assert.equal(getCompactCount(dir), 1);
65
+ assert.equal(getRecallInjected(dir), 4);
66
+ assert.equal(getCacheHitTokensSaved(dir), 500);
67
+ });
68
+ });
@@ -704,6 +704,34 @@ export function bumpDedupStats(deduped: boolean, stateDir: string = getStateDir(
704
704
  if (deduped) incMeta("deduped", 1, stateDir);
705
705
  }
706
706
 
707
+ // --- Live dashboard counters (schemaless meta key/value — NO migration) -----
708
+ // These reuse the private `incMeta` atomically-incrementing integer counter so
709
+ // all cumulative tallies live in the same `meta` table as tokens_saved etc.
710
+
711
+ export function incCompactCount(stateDir: string = getStateDir()): void {
712
+ incMeta("compact_count", 1, stateDir);
713
+ }
714
+
715
+ export function getCompactCount(stateDir: string = getStateDir()): number {
716
+ return getMetaNumber("compact_count", stateDir);
717
+ }
718
+
719
+ export function incRecallInjected(n: number, stateDir: string = getStateDir()): void {
720
+ if (n > 0) incMeta("recall_injected", n, stateDir);
721
+ }
722
+
723
+ export function getRecallInjected(stateDir: string = getStateDir()): number {
724
+ return getMetaNumber("recall_injected", stateDir);
725
+ }
726
+
727
+ export function incCacheHitTokens(delta: number, stateDir: string = getStateDir()): void {
728
+ if (delta > 0) incMeta("cache_hit_tokens_saved", delta, stateDir);
729
+ }
730
+
731
+ export function getCacheHitTokensSaved(stateDir: string = getStateDir()): number {
732
+ return getMetaNumber("cache_hit_tokens_saved", stateDir);
733
+ }
734
+
707
735
  // --- Future-feature foundation (resume sessions / daily log / lessons) -------
708
736
  // Scaffolded tables + minimal helpers so all store data lives in SQLite from
709
737
  // day one. Full UI/recall for these lands in later sprints.
@@ -1,291 +0,0 @@
1
- /**
2
- * openclaw-mega-compact — OpenClaw plugin adapter for the pi-mega-compact engine.
3
- *
4
- * Wires the pi-agnostic Trident engine (src/) into OpenClaw's plugin lifecycle:
5
- * - Registers a CompactionProvider that replaces the built-in summarizeInStages.
6
- * - Exposes `mega_status` and `mega_recall` tools for on-demand inspection.
7
- * - Hooks into `before_compaction` / `after_compaction` for diagnostics.
8
- *
9
- * Design constraints:
10
- * - NO imports from `@earendil-works/pi-coding-agent` or pi-agent-core.
11
- * - The engine core (src/) is pi-agnostic; this file is the sole OpenClaw boundary.
12
- * - No network at runtime — everything is local (stores + extractive summarizer).
13
- */
14
- import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry";
15
- import { compactSession, setDefaultStore, } from "../src/engine.js";
16
- import { recallAndInline } from "../src/recall.js";
17
- import { VectorStore } from "../src/vectorStore.js";
18
- // ---------------------------------------------------------------------------
19
- // Constants
20
- // ---------------------------------------------------------------------------
21
- const PLUGIN_ID = "mega-compact";
22
- const PLUGIN_LABEL = "Mega Compact (Trident)";
23
- /** Default state directory for vector store persistence. */
24
- const STATE_DIR = process.env.MEGA_COMPACT_STATE_DIR ?? undefined;
25
- /** Minimum messages before we bother compacting. */
26
- const MIN_MESSAGES_FOR_COMPACT = 6;
27
- // ---------------------------------------------------------------------------
28
- // Message conversion — OpenClaw unknown[] → EngineMessage[]
29
- // ---------------------------------------------------------------------------
30
- /**
31
- * Best-effort conversion from OpenClaw's opaque message array to our
32
- * EngineMessage shape. OpenClaw messages are typed as `unknown[]` so we
33
- * handle whatever shape comes through gracefully.
34
- */
35
- function toEngineMessages(messages) {
36
- return messages.map((msg) => {
37
- if (!msg || typeof msg !== "object") {
38
- // Primitive fallback — treat as custom text.
39
- return {
40
- role: "custom",
41
- text: String(msg ?? ""),
42
- };
43
- }
44
- const m = msg;
45
- const role = typeof m.role === "string" ? m.role : "custom";
46
- // Normalize role to one of our four engine roles.
47
- let engineRole;
48
- switch (role) {
49
- case "user":
50
- engineRole = "user";
51
- break;
52
- case "assistant":
53
- engineRole = "assistant";
54
- break;
55
- case "tool":
56
- case "function":
57
- engineRole = "tool";
58
- break;
59
- default:
60
- engineRole = "custom";
61
- break;
62
- }
63
- // Extract text content from common message shapes.
64
- const text = typeof m.content === "string"
65
- ? m.content
66
- : typeof m.text === "string"
67
- ? m.text
68
- : Array.isArray(m.content)
69
- ? m.content
70
- .filter((part) => part.type === "text" && typeof part.text === "string")
71
- .map((part) => part.text)
72
- .join("\n")
73
- : "";
74
- // Preserve tool metadata when present.
75
- const toolName = typeof m.name === "string"
76
- ? m.name
77
- : typeof m.toolName === "string"
78
- ? m.toolName
79
- : undefined;
80
- const input = typeof m.input === "string"
81
- ? m.input
82
- : typeof m.arguments === "string"
83
- ? m.arguments
84
- : m.arguments !== undefined
85
- ? JSON.stringify(m.arguments)
86
- : undefined;
87
- const output = typeof m.output === "string"
88
- ? m.output
89
- : engineRole === "tool" && typeof m.content === "string"
90
- ? m.content
91
- : undefined;
92
- return { role: engineRole, text, toolName, input, output };
93
- });
94
- }
95
- // ---------------------------------------------------------------------------
96
- // Compaction provider
97
- // ---------------------------------------------------------------------------
98
- function createCompactionProvider(store) {
99
- return {
100
- id: PLUGIN_ID,
101
- label: PLUGIN_LABEL,
102
- async summarize({ messages, signal, compressionRatio, }) {
103
- // Abort check — bail early if the caller cancelled.
104
- if (signal?.aborted) {
105
- throw new DOMException("Aborted", "AbortError");
106
- }
107
- const engineMessages = toEngineMessages(messages);
108
- // Nothing meaningful to compact.
109
- if (engineMessages.length < MIN_MESSAGES_FOR_COMPACT) {
110
- return "";
111
- }
112
- // Map compression ratio → keepFrom boundary.
113
- // compressionRatio=0.5 means "compact the oldest 50%".
114
- // Default to compacting the oldest half if not specified.
115
- const ratio = compressionRatio ?? 0.5;
116
- const keepFrom = Math.max(MIN_MESSAGES_FOR_COMPACT, Math.floor(engineMessages.length * (1 - ratio)));
117
- // Abort check after conversion (conversion is cheap but check anyway).
118
- if (signal?.aborted) {
119
- throw new DOMException("Aborted", "AbortError");
120
- }
121
- const sessionId = `openclaw-${Date.now()}`;
122
- const input = {
123
- sessionId,
124
- messages: engineMessages,
125
- keepFrom,
126
- };
127
- const result = compactSession(input, store);
128
- if (result.skipped) {
129
- return "";
130
- }
131
- return result.summary;
132
- },
133
- };
134
- }
135
- // ---------------------------------------------------------------------------
136
- // Plugin entry
137
- // ---------------------------------------------------------------------------
138
- export default definePluginEntry({
139
- id: PLUGIN_ID,
140
- name: "Mega Compact",
141
- description: "Layered, local, vector-backed context compressor (Trident engine) for OpenClaw compaction.",
142
- register(api) {
143
- const logger = api.logger;
144
- // Resolve state directory — prefer plugin config override.
145
- const pluginCfg = (api.pluginConfig ?? {});
146
- const stateDir = typeof pluginCfg.stateDir === "string" && pluginCfg.stateDir.length > 0
147
- ? pluginCfg.stateDir
148
- : STATE_DIR;
149
- // Initialize vector store.
150
- let store;
151
- try {
152
- store = new VectorStore({ stateDir });
153
- setDefaultStore(store);
154
- logger.info?.(`${PLUGIN_ID}: vector store initialized (stateDir=${stateDir ?? "default"})`);
155
- }
156
- catch (err) {
157
- logger.error?.(`${PLUGIN_ID}: failed to init vector store:`, err);
158
- return; // Hard bail — no point registering if store is broken.
159
- }
160
- // -----------------------------------------------------------------------
161
- // Register compaction provider
162
- // -----------------------------------------------------------------------
163
- const provider = createCompactionProvider(store);
164
- api.registerCompactionProvider(provider);
165
- logger.info?.(`${PLUGIN_ID}: registered compaction provider "${provider.id}"`);
166
- // -----------------------------------------------------------------------
167
- // Hooks — before / after compaction diagnostics
168
- // -----------------------------------------------------------------------
169
- api.registerHook({
170
- event: "before_compaction",
171
- handler: async (ctx) => {
172
- const msgCount = Array.isArray(ctx?.messages) ? ctx.messages.length : 0;
173
- logger.info?.(`${PLUGIN_ID}: before_compaction — ${msgCount} messages in scope`);
174
- },
175
- });
176
- api.registerHook({
177
- event: "after_compaction",
178
- handler: async (ctx) => {
179
- const summaryLen = typeof ctx?.summary === "string" ? ctx.summary.length : 0;
180
- logger.info?.(`${PLUGIN_ID}: after_compaction — summary ${summaryLen} chars`);
181
- },
182
- });
183
- // -----------------------------------------------------------------------
184
- // Tool: mega_status
185
- // -----------------------------------------------------------------------
186
- api.registerTool({
187
- name: "mega_status",
188
- description: "Show the current status of the mega-compact engine: vector store stats, checkpoint count, and recent compaction activity.",
189
- parameters: {
190
- type: "object",
191
- properties: {
192
- sessionId: {
193
- type: "string",
194
- description: "Optional session ID to scope stats to.",
195
- },
196
- },
197
- additionalProperties: false,
198
- },
199
- handler: async (args) => {
200
- const sessionId = args?.sessionId ?? "global";
201
- try {
202
- const stats = store.stats(sessionId);
203
- const parts = [
204
- `**Mega Compact Status**`,
205
- `Session: ${sessionId}`,
206
- `Checkpoints: ${stats.checkpointCount}`,
207
- `Total tokens saved: ${stats.totalTokenEstimate}`,
208
- `Last checkpoint: ${stats.lastCheckpointId ?? "—"}`,
209
- `Injected count: ${stats.injectedCount}`,
210
- `Dedup hit rate: ${(stats.dedupHitRate * 100).toFixed(0)}%`,
211
- ];
212
- if (stats.lastSummary) {
213
- parts.push(`\nLast summary (truncated):\n ${stats.lastSummary.slice(0, 120).replace(/\n/g, " ")}…`);
214
- }
215
- return { content: [{ type: "text", text: parts.join("\n") }] };
216
- }
217
- catch (err) {
218
- return {
219
- content: [{ type: "text", text: `Error reading mega-compact status: ${err}` }],
220
- isError: true,
221
- };
222
- }
223
- },
224
- });
225
- // -----------------------------------------------------------------------
226
- // Tool: mega_recall
227
- // -----------------------------------------------------------------------
228
- api.registerTool({
229
- name: "mega_recall",
230
- description: "Recall and inline relevant context from the mega-compact vector store for the current session.",
231
- parameters: {
232
- type: "object",
233
- properties: {
234
- sessionId: {
235
- type: "string",
236
- description: "Session ID to recall context for.",
237
- },
238
- query: {
239
- type: "string",
240
- description: "Natural language query for relevant context.",
241
- },
242
- limit: {
243
- type: "number",
244
- description: "Max checkpoints to recall (default 3).",
245
- },
246
- },
247
- required: ["sessionId", "query"],
248
- additionalProperties: false,
249
- },
250
- handler: async (args) => {
251
- const { sessionId, query, limit } = args;
252
- if (!sessionId || !query) {
253
- return {
254
- content: [{ type: "text", text: "Both `sessionId` and `query` are required." }],
255
- isError: true,
256
- };
257
- }
258
- try {
259
- const result = recallAndInline({ sessionId, query, limit: limit ?? 3, source: "command", skipInjected: false }, store);
260
- if (result.toInject.length === 0) {
261
- return {
262
- content: [{ type: "text", text: "No relevant context found in the mega-compact store." }],
263
- };
264
- }
265
- const parts = [
266
- `**Recalled ${result.toInject.length} checkpoint(s):**`,
267
- ...result.report,
268
- "",
269
- "---",
270
- result.block,
271
- ];
272
- return { content: [{ type: "text", text: parts.join("\n") }] };
273
- }
274
- catch (err) {
275
- return {
276
- content: [{ type: "text", text: `Error during mega-recall: ${err}` }],
277
- isError: true,
278
- };
279
- }
280
- },
281
- });
282
- // -----------------------------------------------------------------------
283
- // Cleanup on shutdown
284
- // -----------------------------------------------------------------------
285
- api.on("shutdown", () => {
286
- logger.info?.(`${PLUGIN_ID}: shutting down — clearing default store`);
287
- setDefaultStore(undefined);
288
- });
289
- logger.info?.(`${PLUGIN_ID}: plugin registered (tools: mega_status, mega_recall)`);
290
- },
291
- });
@@ -1,92 +0,0 @@
1
- /**
2
- * minilm.ts — local MiniLM (all-MiniLM-L6-v2) sentence embedder (Sprint 12).
3
- *
4
- * Implements the `Embedder` interface so it drops into the existing VectorStore
5
- * dedup cascade and search with no call-site changes. Inference is 100% local:
6
- * the ONNX model + WordPiece vocab are on-disk artifacts fetched once by
7
- * scripts/setup-minilm.mjs. There is NO network call at runtime (PREVENT-PI-004).
8
- *
9
- * Inputs (dynamic): input_ids, attention_mask, token_type_ids (int64).
10
- * Output: last_hidden_state (batch, seq, 384). We mean-pool over non-padded
11
- * tokens (attention_mask == 1) and L2-normalize → 384-dim unit vector.
12
- *
13
- * The ONNX session + tokenizer are loaded LAZILY on first embed() so the default
14
- * TrigramEmbedder path (and its zero native-init cost) is untouched unless
15
- * MEGACOMPACT_EMBEDDER=minilm is selected.
16
- */
17
- import { join } from "node:path";
18
- import { homedir } from "node:os";
19
- import { existsSync } from "node:fs";
20
- import { l2Normalize, awaitSync } from "./embedder.js";
21
- import { WordPieceTokenizer } from "./wordpiece.js";
22
- export const MINILM_DIM = 384;
23
- export const MINILM_MAX_LEN = 256;
24
- /** Resolve the model directory: MEGACOMPACT_MINILM_DIR > ./models/minilm > ~/.pi … */
25
- function resolveModelDir() {
26
- if (process.env.MEGACOMPACT_MINILM_DIR)
27
- return process.env.MEGACOMPACT_MINILM_DIR;
28
- // Repo-local vendored path (gitignored).
29
- const local = join(process.cwd(), "models", "minilm");
30
- if (existsSync(local))
31
- return local;
32
- return join(homedir(), ".pi", "agent", "extensions", "mega-compact", "models", "minilm");
33
- }
34
- export class MiniLMEmbedder {
35
- dim = MINILM_DIM;
36
- session = null;
37
- tokenizer = null;
38
- modelDir;
39
- loadPromise = null;
40
- constructor(modelDir = resolveModelDir()) {
41
- this.modelDir = modelDir;
42
- }
43
- async ensureLoaded() {
44
- if (this.session && this.tokenizer)
45
- return;
46
- if (this.loadPromise)
47
- return this.loadPromise;
48
- this.loadPromise = (async () => {
49
- const ort = await import("onnxruntime-node");
50
- const modelPath = join(this.modelDir, "model_quantized.onnx");
51
- const vocabPath = join(this.modelDir, "vocab.txt");
52
- if (!existsSync(modelPath) || !existsSync(vocabPath)) {
53
- throw new Error(`MiniLM artifacts missing in ${this.modelDir}. Run: node scripts/setup-minilm.mjs`);
54
- }
55
- // 1 thread is plenty for a single short-region embed and bounds CPU.
56
- this.session = await ort.InferenceSession.create(modelPath, {
57
- executionProviders: ["cpu"],
58
- graphOptimizationLevel: "all",
59
- });
60
- this.tokenizer = WordPieceTokenizer.fromVocabFile(vocabPath);
61
- })();
62
- return this.loadPromise;
63
- }
64
- embed(text) {
65
- awaitSync(this.ensureLoaded());
66
- const enc = this.tokenizer.encode(text, MINILM_MAX_LEN);
67
- const n = enc.inputIds.length;
68
- const BigInt64 = (arr) => arr.map((x) => BigInt(x));
69
- const ort = awaitSync(import("onnxruntime-node"));
70
- const tensors = {
71
- input_ids: new ort.Tensor("int64", BigInt64(enc.inputIds), [1, n]),
72
- attention_mask: new ort.Tensor("int64", BigInt64(enc.attentionMask), [1, n]),
73
- token_type_ids: new ort.Tensor("int64", BigInt64(enc.tokenTypeIds), [1, n]),
74
- };
75
- const out = awaitSync(this.session.run(tensors));
76
- const hidden = out.last_hidden_state.data;
77
- // hidden shape: [1, n, 384]. Mean-pool over non-padded positions.
78
- const pooled = new Array(MINILM_DIM).fill(0);
79
- let count = 0;
80
- for (let i = 0; i < n; i++) {
81
- if (enc.attentionMask[i] === 0)
82
- continue;
83
- const base = i * MINILM_DIM;
84
- for (let d = 0; d < MINILM_DIM; d++)
85
- pooled[d] += hidden[base + d];
86
- count++;
87
- }
88
- if (count === 0)
89
- return l2Normalize(new Array(MINILM_DIM).fill(0));
90
- return l2Normalize(pooled.map((x) => x / count));
91
- }
92
- }