agent-working-memory 0.8.7 → 0.9.0
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/README.md +207 -46
- package/dist/api/routes.js +1 -1
- package/dist/cli/migrate.js +29 -29
- package/dist/cli.js +1 -1
- package/dist/coordination/circuit-breaker.js +23 -23
- package/dist/core/write-pipeline.d.ts.map +1 -1
- package/dist/core/write-pipeline.js +17 -0
- package/dist/core/write-pipeline.js.map +1 -1
- package/dist/engine/activation.d.ts +28 -0
- package/dist/engine/activation.d.ts.map +1 -1
- package/dist/engine/activation.js +341 -11
- package/dist/engine/activation.js.map +1 -1
- package/dist/engine/connections.d.ts +12 -0
- package/dist/engine/connections.d.ts.map +1 -1
- package/dist/engine/connections.js +95 -0
- package/dist/engine/connections.js.map +1 -1
- package/dist/mcp.js +2 -2
- package/dist/storage/pglite-schema.js +143 -143
- package/dist/storage/pglite.js +138 -138
- package/dist/types/engram.d.ts +1 -0
- package/dist/types/engram.d.ts.map +1 -1
- package/package.json +1 -1
- package/src/api/index.ts +3 -3
- package/src/api/routes.ts +1 -1
- package/src/cli/migrate.ts +307 -307
- package/src/cli.ts +1 -1
- package/src/coordination/circuit-breaker.ts +83 -83
- package/src/coordination/failure-modes.ts +50 -50
- package/src/core/decay.ts +63 -63
- package/src/core/embeddings.ts +110 -110
- package/src/core/index.ts +5 -5
- package/src/core/logger.ts +36 -36
- package/src/core/ml-worker-entry.ts +194 -194
- package/src/core/ml-worker.ts +281 -281
- package/src/core/query-expander.ts +122 -122
- package/src/core/reranker.ts +119 -119
- package/src/core/write-pipeline.ts +15 -0
- package/src/engine/activation.ts +328 -11
- package/src/engine/confidence.ts +120 -120
- package/src/engine/connections.ts +94 -0
- package/src/engine/consolidation-scheduler.ts +242 -242
- package/src/engine/eval.ts +102 -102
- package/src/engine/eviction.ts +101 -101
- package/src/engine/index.ts +8 -8
- package/src/engine/retraction.ts +366 -366
- package/src/engine/staging.ts +74 -74
- package/src/mcp.ts +2 -2
- package/src/storage/factory.ts +147 -147
- package/src/storage/index.ts +3 -3
- package/src/storage/pglite-schema.ts +166 -166
- package/src/storage/pglite.ts +1363 -1363
- package/src/storage/store.ts +80 -80
- package/src/types/agent.ts +67 -67
- package/src/types/checkpoint.ts +46 -46
- package/src/types/engram.ts +1 -0
- package/src/types/eval.ts +100 -100
- package/src/types/index.ts +6 -6
package/src/core/logger.ts
CHANGED
|
@@ -1,36 +1,36 @@
|
|
|
1
|
-
// Copyright 2026 Robert Winter / Complete Ideas
|
|
2
|
-
// SPDX-License-Identifier: Apache-2.0
|
|
3
|
-
/**
|
|
4
|
-
* Simple file logger for AWM activity.
|
|
5
|
-
*
|
|
6
|
-
* Appends one line per event to data/awm.log (next to memory.db).
|
|
7
|
-
* Format: ISO timestamp | agent | event | detail
|
|
8
|
-
*
|
|
9
|
-
* Designed for dev pilot observability — know at a glance what's happening.
|
|
10
|
-
*/
|
|
11
|
-
|
|
12
|
-
import { appendFileSync, mkdirSync } from 'node:fs';
|
|
13
|
-
import { dirname, resolve } from 'node:path';
|
|
14
|
-
|
|
15
|
-
let logPath: string | null = null;
|
|
16
|
-
|
|
17
|
-
export function initLogger(dbPath: string): void {
|
|
18
|
-
const dir = dirname(resolve(dbPath));
|
|
19
|
-
mkdirSync(dir, { recursive: true });
|
|
20
|
-
logPath = resolve(dir, 'awm.log');
|
|
21
|
-
}
|
|
22
|
-
|
|
23
|
-
export function log(agentId: string, event: string, detail: string): void {
|
|
24
|
-
if (!logPath) return;
|
|
25
|
-
const ts = new Date().toISOString();
|
|
26
|
-
const line = `${ts} | ${agentId} | ${event} | ${detail}\n`;
|
|
27
|
-
try {
|
|
28
|
-
appendFileSync(logPath, line);
|
|
29
|
-
} catch {
|
|
30
|
-
// Logging should never crash the server
|
|
31
|
-
}
|
|
32
|
-
}
|
|
33
|
-
|
|
34
|
-
export function getLogPath(): string | null {
|
|
35
|
-
return logPath;
|
|
36
|
-
}
|
|
1
|
+
// Copyright 2026 Robert Winter / Complete Ideas
|
|
2
|
+
// SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
/**
|
|
4
|
+
* Simple file logger for AWM activity.
|
|
5
|
+
*
|
|
6
|
+
* Appends one line per event to data/awm.log (next to memory.db).
|
|
7
|
+
* Format: ISO timestamp | agent | event | detail
|
|
8
|
+
*
|
|
9
|
+
* Designed for dev pilot observability — know at a glance what's happening.
|
|
10
|
+
*/
|
|
11
|
+
|
|
12
|
+
import { appendFileSync, mkdirSync } from 'node:fs';
|
|
13
|
+
import { dirname, resolve } from 'node:path';
|
|
14
|
+
|
|
15
|
+
let logPath: string | null = null;
|
|
16
|
+
|
|
17
|
+
export function initLogger(dbPath: string): void {
|
|
18
|
+
const dir = dirname(resolve(dbPath));
|
|
19
|
+
mkdirSync(dir, { recursive: true });
|
|
20
|
+
logPath = resolve(dir, 'awm.log');
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
export function log(agentId: string, event: string, detail: string): void {
|
|
24
|
+
if (!logPath) return;
|
|
25
|
+
const ts = new Date().toISOString();
|
|
26
|
+
const line = `${ts} | ${agentId} | ${event} | ${detail}\n`;
|
|
27
|
+
try {
|
|
28
|
+
appendFileSync(logPath, line);
|
|
29
|
+
} catch {
|
|
30
|
+
// Logging should never crash the server
|
|
31
|
+
}
|
|
32
|
+
}
|
|
33
|
+
|
|
34
|
+
export function getLogPath(): string | null {
|
|
35
|
+
return logPath;
|
|
36
|
+
}
|
|
@@ -1,194 +1,194 @@
|
|
|
1
|
-
// Copyright 2026 Robert Winter / Complete Ideas
|
|
2
|
-
// SPDX-License-Identifier: Apache-2.0
|
|
3
|
-
/**
|
|
4
|
-
* ML worker entry — runs INSIDE a worker_thread.
|
|
5
|
-
*
|
|
6
|
-
* Loaded once per worker, sets up the model for the assigned role
|
|
7
|
-
* (embed | rerank | expand), then handles request messages from the
|
|
8
|
-
* main thread via the parentPort.
|
|
9
|
-
*
|
|
10
|
-
* Protocol:
|
|
11
|
-
* Main thread → worker: { id, op, args } (op matches the worker's role)
|
|
12
|
-
* Worker → main thread: { id, ok: true, result } or { id, ok: false, error }
|
|
13
|
-
* Worker → main thread: { ready: true } (one-time signal after model load)
|
|
14
|
-
* Main thread → worker: { shutdown: true } (drain queue, then terminate)
|
|
15
|
-
*
|
|
16
|
-
* The worker stays loaded — the model lives in memory for the worker's lifetime.
|
|
17
|
-
*/
|
|
18
|
-
|
|
19
|
-
import { parentPort, workerData } from 'node:worker_threads';
|
|
20
|
-
|
|
21
|
-
if (!parentPort) {
|
|
22
|
-
throw new Error('ml-worker-entry: must be loaded as a worker_thread');
|
|
23
|
-
}
|
|
24
|
-
|
|
25
|
-
type WorkerRole = 'embed' | 'rerank' | 'expand';
|
|
26
|
-
const role: WorkerRole = workerData?.role;
|
|
27
|
-
if (role !== 'embed' && role !== 'rerank' && role !== 'expand') {
|
|
28
|
-
throw new Error(`ml-worker-entry: invalid role '${role}'`);
|
|
29
|
-
}
|
|
30
|
-
|
|
31
|
-
// --- Lazy model loaders (each worker loads only its own model) ---
|
|
32
|
-
|
|
33
|
-
let embedderPipeline: any = null;
|
|
34
|
-
let rerankerTokenizer: any = null;
|
|
35
|
-
let rerankerModel: any = null;
|
|
36
|
-
let expanderPipeline: any = null;
|
|
37
|
-
|
|
38
|
-
// Inside worker_threads we must use the WASM ONNX backend, not the native one.
|
|
39
|
-
// onnxruntime-node's native bindings store V8 handles that get invalidated when
|
|
40
|
-
// crossing isolate boundaries — calling from a worker crashes with
|
|
41
|
-
// `v8::HandleScope::CreateHandle()` failures. The WASM backend is V8-safe.
|
|
42
|
-
const WORKER_DEVICE = 'wasm' as const;
|
|
43
|
-
|
|
44
|
-
async function loadEmbedder(): Promise<void> {
|
|
45
|
-
const { pipeline } = await import('@huggingface/transformers');
|
|
46
|
-
const modelId = process.env.AWM_EMBED_MODEL ?? 'Xenova/bge-small-en-v1.5';
|
|
47
|
-
embedderPipeline = await pipeline('feature-extraction', modelId, { dtype: 'fp32', device: WORKER_DEVICE });
|
|
48
|
-
}
|
|
49
|
-
|
|
50
|
-
async function loadReranker(): Promise<void> {
|
|
51
|
-
const { AutoTokenizer, AutoModelForSequenceClassification } = await import('@huggingface/transformers');
|
|
52
|
-
const modelId = process.env.AWM_RERANKER_MODEL || 'Xenova/ms-marco-MiniLM-L-6-v2';
|
|
53
|
-
rerankerTokenizer = await AutoTokenizer.from_pretrained(modelId);
|
|
54
|
-
rerankerModel = await AutoModelForSequenceClassification.from_pretrained(modelId, { dtype: 'fp32', device: WORKER_DEVICE });
|
|
55
|
-
}
|
|
56
|
-
|
|
57
|
-
async function loadExpander(): Promise<void> {
|
|
58
|
-
const { pipeline } = await import('@huggingface/transformers');
|
|
59
|
-
expanderPipeline = await pipeline('text2text-generation', 'Xenova/flan-t5-small', { dtype: 'fp32', device: WORKER_DEVICE });
|
|
60
|
-
}
|
|
61
|
-
|
|
62
|
-
// --- Per-role inference handlers ---
|
|
63
|
-
|
|
64
|
-
function sigmoid(x: number): number {
|
|
65
|
-
return 1 / (1 + Math.exp(-x));
|
|
66
|
-
}
|
|
67
|
-
|
|
68
|
-
async function handleEmbed(args: { texts: string[]; pooling: 'cls' | 'mean'; dimensions: number }): Promise<number[][]> {
|
|
69
|
-
if (!embedderPipeline) throw new Error('embedder not loaded');
|
|
70
|
-
const { texts, pooling, dimensions } = args;
|
|
71
|
-
if (texts.length === 0) return [];
|
|
72
|
-
const result = await embedderPipeline(texts, { pooling, normalize: true });
|
|
73
|
-
const data = result.data as Float32Array;
|
|
74
|
-
const vectors: number[][] = [];
|
|
75
|
-
for (let i = 0; i < texts.length; i++) {
|
|
76
|
-
vectors.push(Array.from(data.slice(i * dimensions, (i + 1) * dimensions)));
|
|
77
|
-
}
|
|
78
|
-
return vectors;
|
|
79
|
-
}
|
|
80
|
-
|
|
81
|
-
interface RerankResult { index: number; score: number; }
|
|
82
|
-
|
|
83
|
-
async function handleRerank(args: { query: string; passages: string[] }): Promise<RerankResult[]> {
|
|
84
|
-
if (!rerankerTokenizer || !rerankerModel) throw new Error('reranker not loaded');
|
|
85
|
-
const { query, passages } = args;
|
|
86
|
-
if (passages.length === 0) return [];
|
|
87
|
-
|
|
88
|
-
// Batch path
|
|
89
|
-
try {
|
|
90
|
-
const queries = passages.map(() => query);
|
|
91
|
-
const inputs = rerankerTokenizer(queries, {
|
|
92
|
-
text_pair: passages,
|
|
93
|
-
padding: true,
|
|
94
|
-
truncation: true,
|
|
95
|
-
return_tensors: 'pt',
|
|
96
|
-
});
|
|
97
|
-
const output = await rerankerModel(inputs);
|
|
98
|
-
const logits = output.logits ?? output.last_hidden_state;
|
|
99
|
-
const data = logits.data as Float32Array | number[];
|
|
100
|
-
const results: RerankResult[] = [];
|
|
101
|
-
for (let i = 0; i < passages.length; i++) {
|
|
102
|
-
const rawLogit = Number(data[i] ?? 0);
|
|
103
|
-
results.push({ index: i, score: sigmoid(rawLogit) });
|
|
104
|
-
}
|
|
105
|
-
results.sort((a, b) => b.score - a.score);
|
|
106
|
-
return results;
|
|
107
|
-
} catch {
|
|
108
|
-
// Per-passage fallback
|
|
109
|
-
const results: RerankResult[] = [];
|
|
110
|
-
for (let i = 0; i < passages.length; i++) {
|
|
111
|
-
try {
|
|
112
|
-
const inputs = rerankerTokenizer(query, {
|
|
113
|
-
text_pair: passages[i],
|
|
114
|
-
padding: true,
|
|
115
|
-
truncation: true,
|
|
116
|
-
return_tensors: 'pt',
|
|
117
|
-
});
|
|
118
|
-
const output = await rerankerModel(inputs);
|
|
119
|
-
const logits = output.logits ?? output.last_hidden_state;
|
|
120
|
-
const rawLogit = logits.data[0] as number;
|
|
121
|
-
results.push({ index: i, score: sigmoid(rawLogit) });
|
|
122
|
-
} catch {
|
|
123
|
-
results.push({ index: i, score: 0 });
|
|
124
|
-
}
|
|
125
|
-
}
|
|
126
|
-
results.sort((a, b) => b.score - a.score);
|
|
127
|
-
return results;
|
|
128
|
-
}
|
|
129
|
-
}
|
|
130
|
-
|
|
131
|
-
async function handleExpand(args: { prompt: string; maxNewTokens: number; noRepeatNgramSize: number }): Promise<string> {
|
|
132
|
-
if (!expanderPipeline) throw new Error('expander not loaded');
|
|
133
|
-
const result = await expanderPipeline(args.prompt, {
|
|
134
|
-
max_new_tokens: args.maxNewTokens,
|
|
135
|
-
no_repeat_ngram_size: args.noRepeatNgramSize,
|
|
136
|
-
});
|
|
137
|
-
const text = Array.isArray(result) ? (result[0] as any)?.generated_text ?? '' : '';
|
|
138
|
-
return String(text).trim();
|
|
139
|
-
}
|
|
140
|
-
|
|
141
|
-
// --- Main loop ---
|
|
142
|
-
|
|
143
|
-
let shuttingDown = false;
|
|
144
|
-
const inflight = new Set<Promise<void>>();
|
|
145
|
-
|
|
146
|
-
async function loadModel(): Promise<void> {
|
|
147
|
-
switch (role) {
|
|
148
|
-
case 'embed': await loadEmbedder(); break;
|
|
149
|
-
case 'rerank': await loadReranker(); break;
|
|
150
|
-
case 'expand': await loadExpander(); break;
|
|
151
|
-
}
|
|
152
|
-
}
|
|
153
|
-
|
|
154
|
-
async function handleMessage(msg: { id: number; op: WorkerRole; args: any }): Promise<void> {
|
|
155
|
-
try {
|
|
156
|
-
let result: unknown;
|
|
157
|
-
switch (msg.op) {
|
|
158
|
-
case 'embed': result = await handleEmbed(msg.args); break;
|
|
159
|
-
case 'rerank': result = await handleRerank(msg.args); break;
|
|
160
|
-
case 'expand': result = await handleExpand(msg.args); break;
|
|
161
|
-
}
|
|
162
|
-
parentPort!.postMessage({ id: msg.id, ok: true, result });
|
|
163
|
-
} catch (err) {
|
|
164
|
-
parentPort!.postMessage({ id: msg.id, ok: false, error: String((err as Error)?.message ?? err) });
|
|
165
|
-
}
|
|
166
|
-
}
|
|
167
|
-
|
|
168
|
-
(async () => {
|
|
169
|
-
try {
|
|
170
|
-
await loadModel();
|
|
171
|
-
parentPort!.postMessage({ ready: true, role });
|
|
172
|
-
} catch (err) {
|
|
173
|
-
parentPort!.postMessage({ ready: false, role, error: String((err as Error)?.message ?? err) });
|
|
174
|
-
process.exit(1);
|
|
175
|
-
}
|
|
176
|
-
})();
|
|
177
|
-
|
|
178
|
-
parentPort.on('message', (msg: any) => {
|
|
179
|
-
if (msg?.shutdown) {
|
|
180
|
-
shuttingDown = true;
|
|
181
|
-
// Wait for in-flight work, then exit
|
|
182
|
-
void Promise.allSettled([...inflight]).then(() => {
|
|
183
|
-
parentPort!.postMessage({ shutdown: 'done' });
|
|
184
|
-
process.exit(0);
|
|
185
|
-
});
|
|
186
|
-
return;
|
|
187
|
-
}
|
|
188
|
-
if (shuttingDown) return;
|
|
189
|
-
if (typeof msg?.id !== 'number' || typeof msg?.op !== 'string') return;
|
|
190
|
-
|
|
191
|
-
const promise = handleMessage(msg);
|
|
192
|
-
inflight.add(promise);
|
|
193
|
-
void promise.finally(() => inflight.delete(promise));
|
|
194
|
-
});
|
|
1
|
+
// Copyright 2026 Robert Winter / Complete Ideas
|
|
2
|
+
// SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
/**
|
|
4
|
+
* ML worker entry — runs INSIDE a worker_thread.
|
|
5
|
+
*
|
|
6
|
+
* Loaded once per worker, sets up the model for the assigned role
|
|
7
|
+
* (embed | rerank | expand), then handles request messages from the
|
|
8
|
+
* main thread via the parentPort.
|
|
9
|
+
*
|
|
10
|
+
* Protocol:
|
|
11
|
+
* Main thread → worker: { id, op, args } (op matches the worker's role)
|
|
12
|
+
* Worker → main thread: { id, ok: true, result } or { id, ok: false, error }
|
|
13
|
+
* Worker → main thread: { ready: true } (one-time signal after model load)
|
|
14
|
+
* Main thread → worker: { shutdown: true } (drain queue, then terminate)
|
|
15
|
+
*
|
|
16
|
+
* The worker stays loaded — the model lives in memory for the worker's lifetime.
|
|
17
|
+
*/
|
|
18
|
+
|
|
19
|
+
import { parentPort, workerData } from 'node:worker_threads';
|
|
20
|
+
|
|
21
|
+
if (!parentPort) {
|
|
22
|
+
throw new Error('ml-worker-entry: must be loaded as a worker_thread');
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
type WorkerRole = 'embed' | 'rerank' | 'expand';
|
|
26
|
+
const role: WorkerRole = workerData?.role;
|
|
27
|
+
if (role !== 'embed' && role !== 'rerank' && role !== 'expand') {
|
|
28
|
+
throw new Error(`ml-worker-entry: invalid role '${role}'`);
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
// --- Lazy model loaders (each worker loads only its own model) ---
|
|
32
|
+
|
|
33
|
+
let embedderPipeline: any = null;
|
|
34
|
+
let rerankerTokenizer: any = null;
|
|
35
|
+
let rerankerModel: any = null;
|
|
36
|
+
let expanderPipeline: any = null;
|
|
37
|
+
|
|
38
|
+
// Inside worker_threads we must use the WASM ONNX backend, not the native one.
|
|
39
|
+
// onnxruntime-node's native bindings store V8 handles that get invalidated when
|
|
40
|
+
// crossing isolate boundaries — calling from a worker crashes with
|
|
41
|
+
// `v8::HandleScope::CreateHandle()` failures. The WASM backend is V8-safe.
|
|
42
|
+
const WORKER_DEVICE = 'wasm' as const;
|
|
43
|
+
|
|
44
|
+
async function loadEmbedder(): Promise<void> {
|
|
45
|
+
const { pipeline } = await import('@huggingface/transformers');
|
|
46
|
+
const modelId = process.env.AWM_EMBED_MODEL ?? 'Xenova/bge-small-en-v1.5';
|
|
47
|
+
embedderPipeline = await pipeline('feature-extraction', modelId, { dtype: 'fp32', device: WORKER_DEVICE });
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
async function loadReranker(): Promise<void> {
|
|
51
|
+
const { AutoTokenizer, AutoModelForSequenceClassification } = await import('@huggingface/transformers');
|
|
52
|
+
const modelId = process.env.AWM_RERANKER_MODEL || 'Xenova/ms-marco-MiniLM-L-6-v2';
|
|
53
|
+
rerankerTokenizer = await AutoTokenizer.from_pretrained(modelId);
|
|
54
|
+
rerankerModel = await AutoModelForSequenceClassification.from_pretrained(modelId, { dtype: 'fp32', device: WORKER_DEVICE });
|
|
55
|
+
}
|
|
56
|
+
|
|
57
|
+
async function loadExpander(): Promise<void> {
|
|
58
|
+
const { pipeline } = await import('@huggingface/transformers');
|
|
59
|
+
expanderPipeline = await pipeline('text2text-generation', 'Xenova/flan-t5-small', { dtype: 'fp32', device: WORKER_DEVICE });
|
|
60
|
+
}
|
|
61
|
+
|
|
62
|
+
// --- Per-role inference handlers ---
|
|
63
|
+
|
|
64
|
+
function sigmoid(x: number): number {
|
|
65
|
+
return 1 / (1 + Math.exp(-x));
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
async function handleEmbed(args: { texts: string[]; pooling: 'cls' | 'mean'; dimensions: number }): Promise<number[][]> {
|
|
69
|
+
if (!embedderPipeline) throw new Error('embedder not loaded');
|
|
70
|
+
const { texts, pooling, dimensions } = args;
|
|
71
|
+
if (texts.length === 0) return [];
|
|
72
|
+
const result = await embedderPipeline(texts, { pooling, normalize: true });
|
|
73
|
+
const data = result.data as Float32Array;
|
|
74
|
+
const vectors: number[][] = [];
|
|
75
|
+
for (let i = 0; i < texts.length; i++) {
|
|
76
|
+
vectors.push(Array.from(data.slice(i * dimensions, (i + 1) * dimensions)));
|
|
77
|
+
}
|
|
78
|
+
return vectors;
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
interface RerankResult { index: number; score: number; }
|
|
82
|
+
|
|
83
|
+
async function handleRerank(args: { query: string; passages: string[] }): Promise<RerankResult[]> {
|
|
84
|
+
if (!rerankerTokenizer || !rerankerModel) throw new Error('reranker not loaded');
|
|
85
|
+
const { query, passages } = args;
|
|
86
|
+
if (passages.length === 0) return [];
|
|
87
|
+
|
|
88
|
+
// Batch path
|
|
89
|
+
try {
|
|
90
|
+
const queries = passages.map(() => query);
|
|
91
|
+
const inputs = rerankerTokenizer(queries, {
|
|
92
|
+
text_pair: passages,
|
|
93
|
+
padding: true,
|
|
94
|
+
truncation: true,
|
|
95
|
+
return_tensors: 'pt',
|
|
96
|
+
});
|
|
97
|
+
const output = await rerankerModel(inputs);
|
|
98
|
+
const logits = output.logits ?? output.last_hidden_state;
|
|
99
|
+
const data = logits.data as Float32Array | number[];
|
|
100
|
+
const results: RerankResult[] = [];
|
|
101
|
+
for (let i = 0; i < passages.length; i++) {
|
|
102
|
+
const rawLogit = Number(data[i] ?? 0);
|
|
103
|
+
results.push({ index: i, score: sigmoid(rawLogit) });
|
|
104
|
+
}
|
|
105
|
+
results.sort((a, b) => b.score - a.score);
|
|
106
|
+
return results;
|
|
107
|
+
} catch {
|
|
108
|
+
// Per-passage fallback
|
|
109
|
+
const results: RerankResult[] = [];
|
|
110
|
+
for (let i = 0; i < passages.length; i++) {
|
|
111
|
+
try {
|
|
112
|
+
const inputs = rerankerTokenizer(query, {
|
|
113
|
+
text_pair: passages[i],
|
|
114
|
+
padding: true,
|
|
115
|
+
truncation: true,
|
|
116
|
+
return_tensors: 'pt',
|
|
117
|
+
});
|
|
118
|
+
const output = await rerankerModel(inputs);
|
|
119
|
+
const logits = output.logits ?? output.last_hidden_state;
|
|
120
|
+
const rawLogit = logits.data[0] as number;
|
|
121
|
+
results.push({ index: i, score: sigmoid(rawLogit) });
|
|
122
|
+
} catch {
|
|
123
|
+
results.push({ index: i, score: 0 });
|
|
124
|
+
}
|
|
125
|
+
}
|
|
126
|
+
results.sort((a, b) => b.score - a.score);
|
|
127
|
+
return results;
|
|
128
|
+
}
|
|
129
|
+
}
|
|
130
|
+
|
|
131
|
+
async function handleExpand(args: { prompt: string; maxNewTokens: number; noRepeatNgramSize: number }): Promise<string> {
|
|
132
|
+
if (!expanderPipeline) throw new Error('expander not loaded');
|
|
133
|
+
const result = await expanderPipeline(args.prompt, {
|
|
134
|
+
max_new_tokens: args.maxNewTokens,
|
|
135
|
+
no_repeat_ngram_size: args.noRepeatNgramSize,
|
|
136
|
+
});
|
|
137
|
+
const text = Array.isArray(result) ? (result[0] as any)?.generated_text ?? '' : '';
|
|
138
|
+
return String(text).trim();
|
|
139
|
+
}
|
|
140
|
+
|
|
141
|
+
// --- Main loop ---
|
|
142
|
+
|
|
143
|
+
let shuttingDown = false;
|
|
144
|
+
const inflight = new Set<Promise<void>>();
|
|
145
|
+
|
|
146
|
+
async function loadModel(): Promise<void> {
|
|
147
|
+
switch (role) {
|
|
148
|
+
case 'embed': await loadEmbedder(); break;
|
|
149
|
+
case 'rerank': await loadReranker(); break;
|
|
150
|
+
case 'expand': await loadExpander(); break;
|
|
151
|
+
}
|
|
152
|
+
}
|
|
153
|
+
|
|
154
|
+
async function handleMessage(msg: { id: number; op: WorkerRole; args: any }): Promise<void> {
|
|
155
|
+
try {
|
|
156
|
+
let result: unknown;
|
|
157
|
+
switch (msg.op) {
|
|
158
|
+
case 'embed': result = await handleEmbed(msg.args); break;
|
|
159
|
+
case 'rerank': result = await handleRerank(msg.args); break;
|
|
160
|
+
case 'expand': result = await handleExpand(msg.args); break;
|
|
161
|
+
}
|
|
162
|
+
parentPort!.postMessage({ id: msg.id, ok: true, result });
|
|
163
|
+
} catch (err) {
|
|
164
|
+
parentPort!.postMessage({ id: msg.id, ok: false, error: String((err as Error)?.message ?? err) });
|
|
165
|
+
}
|
|
166
|
+
}
|
|
167
|
+
|
|
168
|
+
(async () => {
|
|
169
|
+
try {
|
|
170
|
+
await loadModel();
|
|
171
|
+
parentPort!.postMessage({ ready: true, role });
|
|
172
|
+
} catch (err) {
|
|
173
|
+
parentPort!.postMessage({ ready: false, role, error: String((err as Error)?.message ?? err) });
|
|
174
|
+
process.exit(1);
|
|
175
|
+
}
|
|
176
|
+
})();
|
|
177
|
+
|
|
178
|
+
parentPort.on('message', (msg: any) => {
|
|
179
|
+
if (msg?.shutdown) {
|
|
180
|
+
shuttingDown = true;
|
|
181
|
+
// Wait for in-flight work, then exit
|
|
182
|
+
void Promise.allSettled([...inflight]).then(() => {
|
|
183
|
+
parentPort!.postMessage({ shutdown: 'done' });
|
|
184
|
+
process.exit(0);
|
|
185
|
+
});
|
|
186
|
+
return;
|
|
187
|
+
}
|
|
188
|
+
if (shuttingDown) return;
|
|
189
|
+
if (typeof msg?.id !== 'number' || typeof msg?.op !== 'string') return;
|
|
190
|
+
|
|
191
|
+
const promise = handleMessage(msg);
|
|
192
|
+
inflight.add(promise);
|
|
193
|
+
void promise.finally(() => inflight.delete(promise));
|
|
194
|
+
});
|