agent-working-memory 0.7.17 → 0.8.6

package/src/core/embeddings.ts

@@ -1,93 +1,110 @@
-// Copyright 2026 Robert Winter / Complete Ideas
-// SPDX-License-Identifier: Apache-2.0
-/**
- * Embedding Engine — local vector embeddings via transformers.js
- *
- * Default: bge-small-en-v1.5 (384 dimensions, ~90MB, MTEB retrieval-optimized).
- * Better short-text similarity than MiniLM for agent memory concepts.
- * Configurable via AWM_EMBED_MODEL env var.
- * Model is downloaded once on first use and cached locally.
- *
- * Singleton pattern — call getEmbedder() to get the shared instance.
- *
- * NOTE: Changing the model invalidates existing embeddings.
- * Set AWM_EMBED_MODEL=Xenova/all-MiniLM-L6-v2 for backward compatibility.
- */
-
-import { pipeline, type FeatureExtractionPipeline } from '@huggingface/transformers';
-
-const MODEL_ID = process.env.AWM_EMBED_MODEL ?? 'Xenova/bge-small-en-v1.5';
-const DIMENSIONS = parseInt(process.env.AWM_EMBED_DIMS ?? '384', 10);
-const POOLING = (process.env.AWM_EMBED_POOLING ?? 'mean') as 'cls' | 'mean';
-
-let instance: FeatureExtractionPipeline | null = null;
-let initPromise: Promise<FeatureExtractionPipeline> | null = null;
-
-/**
- * Get or initialize the embedding pipeline (singleton).
- * First call downloads the model (~22MB), subsequent calls are instant.
- */
-export async function getEmbedder(): Promise<FeatureExtractionPipeline> {
-  if (instance) return instance;
-  if (initPromise) return initPromise;
-
-  initPromise = pipeline('feature-extraction', MODEL_ID, {
-    dtype: 'fp32',
-  }).then(pipe => {
-    instance = pipe;
-    console.log(`Embedding model loaded: ${MODEL_ID} (${DIMENSIONS}d)`);
-    return pipe;
-  });
-
-  return initPromise;
-}
-
-/**
- * Generate an embedding vector for a text string.
- * Returns a normalized float32 array of length DIMENSIONS.
- */
-export async function embed(text: string): Promise<number[]> {
-  const embedder = await getEmbedder();
-  const result = await embedder(text, { pooling: POOLING, normalize: true });
-  // result is a Tensor — extract the data
-  return Array.from(result.data as Float32Array).slice(0, DIMENSIONS);
-}
-
-/** Get the current embedding model ID (for version tracking in stored embeddings) */
-export function getModelId(): string {
-  return MODEL_ID;
-}
-
-/**
- * Generate embeddings for multiple texts in a batch.
- * More efficient than calling embed() in a loop.
- */
-export async function embedBatch(texts: string[]): Promise<number[][]> {
-  if (texts.length === 0) return [];
-  const embedder = await getEmbedder();
-  const result = await embedder(texts, { pooling: POOLING, normalize: true });
-  const data = result.data as Float32Array;
-
-  const vectors: number[][] = [];
-  for (let i = 0; i < texts.length; i++) {
-    vectors.push(Array.from(data.slice(i * DIMENSIONS, (i + 1) * DIMENSIONS)));
-  }
-  return vectors;
-}
-
-/**
- * Cosine similarity between two normalized vectors.
- * Since vectors are pre-normalized, this is just the dot product.
- */
-export function cosineSimilarity(a: number[], b: number[]): number {
-  if (a.length !== b.length || a.length === 0) return 0;
-  let dot = 0;
-  for (let i = 0; i < a.length; i++) {
-    dot += a[i] * b[i];
-  }
-  // Clamp to [-1, 1] to handle floating point drift
-  return Math.max(-1, Math.min(1, dot));
-}
-
-/** Vector dimensions for this model */
-export const EMBEDDING_DIMENSIONS = DIMENSIONS;
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * Embedding Engine - vector embeddings via the ML worker pool.
+ *
+ * Default model: bge-small-en-v1.5 (384 dimensions, ~90MB, MTEB retrieval-optimized).
+ * Configurable via AWM_EMBED_MODEL env var.
+ *
+ * AWM 0.8.x: inference dispatches through ml-worker.ts. The worker_threads
+ * path was planned but reverted to in-process because onnxruntime-node's
+ * native bindings store V8 handles that don't cross isolate boundaries
+ * safely — see ml-worker.ts for the full status. The dispatch abstraction
+ * is preserved for a future child_process or HTTP sidecar pool.
+ * `AWM_ML_INPROCESS=1` is honored as a no-op (in-process is now the default).
+ *
+ * NOTE: Changing the model invalidates existing embeddings.
+ * Set AWM_EMBED_MODEL=Xenova/all-MiniLM-L6-v2 for backward compatibility.
+ */
+
+import { pipeline, type FeatureExtractionPipeline } from '@huggingface/transformers';
+import { dispatchEmbed, registerInProcessHandlers } from './ml-worker.js';
+
+const MODEL_ID = process.env.AWM_EMBED_MODEL ?? 'Xenova/bge-small-en-v1.5';
+const DIMENSIONS = parseInt(process.env.AWM_EMBED_DIMS ?? '384', 10);
+const POOLING = (process.env.AWM_EMBED_POOLING ?? 'mean') as 'cls' | 'mean';
+
+// --- In-process fallback (used by tests and crash recovery) ---
+
+let inProcessInstance: FeatureExtractionPipeline | null = null;
+let inProcessInitPromise: Promise<FeatureExtractionPipeline> | null = null;
+
+async function loadInProcess(): Promise<FeatureExtractionPipeline> {
+  if (inProcessInstance) return inProcessInstance;
+  if (inProcessInitPromise) return inProcessInitPromise;
+  inProcessInitPromise = pipeline('feature-extraction', MODEL_ID, { dtype: 'fp32' }).then(pipe => {
+    inProcessInstance = pipe;
+    console.log(`Embedding model loaded in-process: ${MODEL_ID} (${DIMENSIONS}d)`);
+    return pipe;
+  });
+  return inProcessInitPromise;
+}
+
+async function inProcessEmbed(args: { texts: string[]; pooling: 'cls' | 'mean'; dimensions: number }): Promise<number[][]> {
+  const { texts, pooling, dimensions } = args;
+  if (texts.length === 0) return [];
+  const embedder = await loadInProcess();
+  const result = await embedder(texts, { pooling, normalize: true });
+  const data = result.data as Float32Array;
+  const vectors: number[][] = [];
+  for (let i = 0; i < texts.length; i++) {
+    vectors.push(Array.from(data.slice(i * dimensions, (i + 1) * dimensions)));
+  }
+  return vectors;
+}
+
+// Register the in-process handler with the pool (used in test mode and as fallback)
+registerInProcessHandlers({ embed: inProcessEmbed });
+
+// --- Public API ---
+
+/**
+ * Get or initialize the embedding pipeline (singleton).
+ * Kept for backwards compat — returns the in-process pipeline only.
+ * Most consumers should use embed() / embedBatch() which dispatch
+ * to the worker pool by default.
+ */
+export async function getEmbedder(): Promise<FeatureExtractionPipeline> {
+  return loadInProcess();
+}
+
+/**
+ * Generate an embedding vector for a text string.
+ * Dispatches to the worker pool (or in-process fallback).
+ */
+export async function embed(text: string): Promise<number[]> {
+  const vectors = await dispatchEmbed({ texts: [text], pooling: POOLING, dimensions: DIMENSIONS });
+  return vectors[0] ?? new Array(DIMENSIONS).fill(0);
+}
+
+/**
+ * Generate embeddings for multiple texts in a batch.
+ * More efficient than calling embed() in a loop — the worker batches the
+ * tokenization + forward pass.
+ */
+export async function embedBatch(texts: string[]): Promise<number[][]> {
+  if (texts.length === 0) return [];
+  return dispatchEmbed({ texts, pooling: POOLING, dimensions: DIMENSIONS });
+}
+
+/** Get the current embedding model ID (for version tracking in stored embeddings) */
+export function getModelId(): string {
+  return MODEL_ID;
+}
+
+/**
+ * Cosine similarity between two normalized vectors.
+ * Since vectors are pre-normalized, this is just the dot product.
+ */
+export function cosineSimilarity(a: number[], b: number[]): number {
+  if (a.length !== b.length || a.length === 0) return 0;
+  let dot = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+  }
+  // Clamp to [-1, 1] to handle floating point drift
+  return Math.max(-1, Math.min(1, dot));
+}
+
+/** Vector dimensions for this model */
+export const EMBEDDING_DIMENSIONS = DIMENSIONS;

package/src/core/index.ts

@@ -1,5 +1,5 @@
-// Copyright 2026 Robert Winter / Complete Ideas
-// SPDX-License-Identifier: Apache-2.0
-export * from './decay.js';
-export * from './hebbian.js';
-export * from './salience.js';
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+export * from './decay.js';
+export * from './hebbian.js';
+export * from './salience.js';

package/src/core/logger.ts

@@ -1,36 +1,36 @@
-// Copyright 2026 Robert Winter / Complete Ideas
-// SPDX-License-Identifier: Apache-2.0
-/**
- * Simple file logger for AWM activity.
- *
- * Appends one line per event to data/awm.log (next to memory.db).
- * Format: ISO timestamp | agent | event | detail
- *
- * Designed for dev pilot observability — know at a glance what's happening.
- */
-
-import { appendFileSync, mkdirSync } from 'node:fs';
-import { dirname, resolve } from 'node:path';
-
-let logPath: string | null = null;
-
-export function initLogger(dbPath: string): void {
-  const dir = dirname(resolve(dbPath));
-  mkdirSync(dir, { recursive: true });
-  logPath = resolve(dir, 'awm.log');
-}
-
-export function log(agentId: string, event: string, detail: string): void {
-  if (!logPath) return;
-  const ts = new Date().toISOString();
-  const line = `${ts} | ${agentId} | ${event} | ${detail}\n`;
-  try {
-    appendFileSync(logPath, line);
-  } catch {
-    // Logging should never crash the server
-  }
-}
-
-export function getLogPath(): string | null {
-  return logPath;
-}
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * Simple file logger for AWM activity.
+ *
+ * Appends one line per event to data/awm.log (next to memory.db).
+ * Format: ISO timestamp | agent | event | detail
+ *
+ * Designed for dev pilot observability — know at a glance what's happening.
+ */
+
+import { appendFileSync, mkdirSync } from 'node:fs';
+import { dirname, resolve } from 'node:path';
+
+let logPath: string | null = null;
+
+export function initLogger(dbPath: string): void {
+  const dir = dirname(resolve(dbPath));
+  mkdirSync(dir, { recursive: true });
+  logPath = resolve(dir, 'awm.log');
+}
+
+export function log(agentId: string, event: string, detail: string): void {
+  if (!logPath) return;
+  const ts = new Date().toISOString();
+  const line = `${ts} | ${agentId} | ${event} | ${detail}\n`;
+  try {
+    appendFileSync(logPath, line);
+  } catch {
+    // Logging should never crash the server
+  }
+}
+
+export function getLogPath(): string | null {
+  return logPath;
+}

package/src/core/ml-worker-entry.ts

@@ -0,0 +1,194 @@
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * ML worker entry — runs INSIDE a worker_thread.
+ *
+ * Loaded once per worker, sets up the model for the assigned role
+ * (embed | rerank | expand), then handles request messages from the
+ * main thread via the parentPort.
+ *
+ * Protocol:
+ *   Main thread → worker:  { id, op, args }    (op matches the worker's role)
+ *   Worker → main thread:  { id, ok: true, result } or { id, ok: false, error }
+ *   Worker → main thread:  { ready: true }     (one-time signal after model load)
+ *   Main thread → worker:  { shutdown: true }  (drain queue, then terminate)
+ *
+ * The worker stays loaded — the model lives in memory for the worker's lifetime.
+ */
+
+import { parentPort, workerData } from 'node:worker_threads';
+
+if (!parentPort) {
+  throw new Error('ml-worker-entry: must be loaded as a worker_thread');
+}
+
+type WorkerRole = 'embed' | 'rerank' | 'expand';
+const role: WorkerRole = workerData?.role;
+if (role !== 'embed' && role !== 'rerank' && role !== 'expand') {
+  throw new Error(`ml-worker-entry: invalid role '${role}'`);
+}
+
+// --- Lazy model loaders (each worker loads only its own model) ---
+
+let embedderPipeline: any = null;
+let rerankerTokenizer: any = null;
+let rerankerModel: any = null;
+let expanderPipeline: any = null;
+
+// Inside worker_threads we must use the WASM ONNX backend, not the native one.
+// onnxruntime-node's native bindings store V8 handles that get invalidated when
+// crossing isolate boundaries — calling from a worker crashes with
+// `v8::HandleScope::CreateHandle()` failures. The WASM backend is V8-safe.
+const WORKER_DEVICE = 'wasm' as const;
+
+async function loadEmbedder(): Promise<void> {
+  const { pipeline } = await import('@huggingface/transformers');
+  const modelId = process.env.AWM_EMBED_MODEL ?? 'Xenova/bge-small-en-v1.5';
+  embedderPipeline = await pipeline('feature-extraction', modelId, { dtype: 'fp32', device: WORKER_DEVICE });
+}
+
+async function loadReranker(): Promise<void> {
+  const { AutoTokenizer, AutoModelForSequenceClassification } = await import('@huggingface/transformers');
+  const modelId = process.env.AWM_RERANKER_MODEL || 'Xenova/ms-marco-MiniLM-L-6-v2';
+  rerankerTokenizer = await AutoTokenizer.from_pretrained(modelId);
+  rerankerModel = await AutoModelForSequenceClassification.from_pretrained(modelId, { dtype: 'fp32', device: WORKER_DEVICE });
+}
+
+async function loadExpander(): Promise<void> {
+  const { pipeline } = await import('@huggingface/transformers');
+  expanderPipeline = await pipeline('text2text-generation', 'Xenova/flan-t5-small', { dtype: 'fp32', device: WORKER_DEVICE });
+}
+
+// --- Per-role inference handlers ---
+
+function sigmoid(x: number): number {
+  return 1 / (1 + Math.exp(-x));
+}
+
+async function handleEmbed(args: { texts: string[]; pooling: 'cls' | 'mean'; dimensions: number }): Promise<number[][]> {
+  if (!embedderPipeline) throw new Error('embedder not loaded');
+  const { texts, pooling, dimensions } = args;
+  if (texts.length === 0) return [];
+  const result = await embedderPipeline(texts, { pooling, normalize: true });
+  const data = result.data as Float32Array;
+  const vectors: number[][] = [];
+  for (let i = 0; i < texts.length; i++) {
+    vectors.push(Array.from(data.slice(i * dimensions, (i + 1) * dimensions)));
+  }
+  return vectors;
+}
+
+interface RerankResult { index: number; score: number; }
+
+async function handleRerank(args: { query: string; passages: string[] }): Promise<RerankResult[]> {
+  if (!rerankerTokenizer || !rerankerModel) throw new Error('reranker not loaded');
+  const { query, passages } = args;
+  if (passages.length === 0) return [];
+
+  // Batch path
+  try {
+    const queries = passages.map(() => query);
+    const inputs = rerankerTokenizer(queries, {
+      text_pair: passages,
+      padding: true,
+      truncation: true,
+      return_tensors: 'pt',
+    });
+    const output = await rerankerModel(inputs);
+    const logits = output.logits ?? output.last_hidden_state;
+    const data = logits.data as Float32Array | number[];
+    const results: RerankResult[] = [];
+    for (let i = 0; i < passages.length; i++) {
+      const rawLogit = Number(data[i] ?? 0);
+      results.push({ index: i, score: sigmoid(rawLogit) });
+    }
+    results.sort((a, b) => b.score - a.score);
+    return results;
+  } catch {
+    // Per-passage fallback
+    const results: RerankResult[] = [];
+    for (let i = 0; i < passages.length; i++) {
+      try {
+        const inputs = rerankerTokenizer(query, {
+          text_pair: passages[i],
+          padding: true,
+          truncation: true,
+          return_tensors: 'pt',
+        });
+        const output = await rerankerModel(inputs);
+        const logits = output.logits ?? output.last_hidden_state;
+        const rawLogit = logits.data[0] as number;
+        results.push({ index: i, score: sigmoid(rawLogit) });
+      } catch {
+        results.push({ index: i, score: 0 });
+      }
+    }
+    results.sort((a, b) => b.score - a.score);
+    return results;
+  }
+}
+
+async function handleExpand(args: { prompt: string; maxNewTokens: number; noRepeatNgramSize: number }): Promise<string> {
+  if (!expanderPipeline) throw new Error('expander not loaded');
+  const result = await expanderPipeline(args.prompt, {
+    max_new_tokens: args.maxNewTokens,
+    no_repeat_ngram_size: args.noRepeatNgramSize,
+  });
+  const text = Array.isArray(result) ? (result[0] as any)?.generated_text ?? '' : '';
+  return String(text).trim();
+}
+
+// --- Main loop ---
+
+let shuttingDown = false;
+const inflight = new Set<Promise<void>>();
+
+async function loadModel(): Promise<void> {
+  switch (role) {
+    case 'embed': await loadEmbedder(); break;
+    case 'rerank': await loadReranker(); break;
+    case 'expand': await loadExpander(); break;
+  }
+}
+
+async function handleMessage(msg: { id: number; op: WorkerRole; args: any }): Promise<void> {
+  try {
+    let result: unknown;
+    switch (msg.op) {
+      case 'embed':  result = await handleEmbed(msg.args); break;
+      case 'rerank': result = await handleRerank(msg.args); break;
+      case 'expand': result = await handleExpand(msg.args); break;
+    }
+    parentPort!.postMessage({ id: msg.id, ok: true, result });
+  } catch (err) {
+    parentPort!.postMessage({ id: msg.id, ok: false, error: String((err as Error)?.message ?? err) });
+  }
+}
+
+(async () => {
+  try {
+    await loadModel();
+    parentPort!.postMessage({ ready: true, role });
+  } catch (err) {
+    parentPort!.postMessage({ ready: false, role, error: String((err as Error)?.message ?? err) });
+    process.exit(1);
+  }
+})();
+
+parentPort.on('message', (msg: any) => {
+  if (msg?.shutdown) {
+    shuttingDown = true;
+    // Wait for in-flight work, then exit
+    void Promise.allSettled([...inflight]).then(() => {
+      parentPort!.postMessage({ shutdown: 'done' });
+      process.exit(0);
+    });
+    return;
+  }
+  if (shuttingDown) return;
+  if (typeof msg?.id !== 'number' || typeof msg?.op !== 'string') return;
+
+  const promise = handleMessage(msg);
+  inflight.add(promise);
+  void promise.finally(() => inflight.delete(promise));
+});