Implement vector-aware memory retrieval and prompt formatter

src/services/memory/formatter.js

@@ -0,0 +1,106 @@
+const DEFAULT_TOKEN_BUDGET = 1000;
+const DEFAULT_MAX_SUMMARY_CHARS = 240;
+const DEFAULT_MAX_CONTEXT_CHARS = 320;
+
+function clampString(value, maxChars) {
+  if (typeof value !== 'string') return '';
+  if (!Number.isFinite(maxChars) || maxChars <= 0) return '';
+  if (value.length <= maxChars) return value;
+  return `${value.slice(0, Math.max(0, maxChars - 1))}…`;
+}
+
+/**
+ * Fast heuristic (approx 4 chars/token) suitable for guardrail budgeting.
+ * @param {string} text
+ */
+export function estimateTokenCount(text) {
+  const normalized = String(text ?? '');
+  if (normalized.length === 0) return 0;
+  return Math.ceil(normalized.length / 4);
+}
+
+/**
+ * @param {import('./schema.js').MemoryRecord} memory
+ * @param {{ maxSummaryChars?: number, maxContextChars?: number }} [options]
+ */
+export function formatMemoryBlock(memory, options = {}) {
+  const summary = clampString(memory.summary ?? '', options.maxSummaryChars ?? DEFAULT_MAX_SUMMARY_CHARS);
+  const shortContext = clampString(memory.content ?? '', options.maxContextChars ?? DEFAULT_MAX_CONTEXT_CHARS);
+  const timestamp = new Date(memory.created_at).toISOString();
+  const tags = Array.isArray(memory.tags) ? memory.tags.join(',') : '';
+
+  return {
+    id: memory.id,
+    timestamp,
+    importance: Number(memory.importance ?? 0).toFixed(2),
+    summary,
+    short_context: shortContext,
+    tags,
+  };
+}
+
+
+function fitBlockToBudget(memory, options, maxTokens) {
+  let maxSummaryChars = options.maxSummaryChars ?? DEFAULT_MAX_SUMMARY_CHARS;
+  let maxContextChars = options.maxContextChars ?? DEFAULT_MAX_CONTEXT_CHARS;
+
+  while (maxSummaryChars >= 8 && maxContextChars >= 8) {
+    const block = formatMemoryBlock(memory, { maxSummaryChars, maxContextChars });
+    const line = blockToLine(block);
+    if (estimateTokenCount(line) <= maxTokens) {
+      return { block, line };
+    }
+
+    if (maxContextChars >= maxSummaryChars) {
+      maxContextChars = Math.max(8, Math.floor(maxContextChars * 0.7));
+    } else {
+      maxSummaryChars = Math.max(8, Math.floor(maxSummaryChars * 0.7));
+    }
+
+    if (maxSummaryChars === 8 && maxContextChars === 8) break;
+  }
+
+  const block = formatMemoryBlock(memory, { maxSummaryChars: 8, maxContextChars: 8 });
+  return { block, line: blockToLine(block) };
+}
+function blockToLine(block) {
+  return `[${block.id}, ${block.timestamp}, ${block.importance}, ${block.summary}, ${block.short_context}, ${block.tags}]`;
+}
+
+/**
+ * @param {import('./schema.js').MemoryRecord[]} memories
+ * @param {{ k?: number, tokenBudget?: number, maxSummaryChars?: number, maxContextChars?: number }} [options]
+ */
+export function formatMemoriesForPrompt(memories, options = {}) {
+  const k = Number.isFinite(options.k) && options.k > 0 ? Math.floor(options.k) : memories.length;
+  const tokenBudget = Number.isFinite(options.tokenBudget) && options.tokenBudget > 0
+    ? Math.floor(options.tokenBudget)
+    : DEFAULT_TOKEN_BUDGET;
+
+  const blocks = [];
+  let usedTokens = 0;
+
+  for (const memory of memories.slice(0, k)) {
+    const remainingTokens = tokenBudget - usedTokens;
+    if (remainingTokens <= 0) break;
+
+    const fitted = fitBlockToBudget(memory, options, remainingTokens);
+    const block = fitted.block;
+    const line = fitted.line;
+    const lineTokens = estimateTokenCount(line);
+
+    if (usedTokens + lineTokens > tokenBudget) {
+      break;
+    }
+
+    blocks.push(block);
+    usedTokens += lineTokens;
+  }
+
+  return {
+    blocks,
+    text: blocks.map(blockToLine).join('\n'),
+    usedTokens,
+    truncated: blocks.length < Math.min(k, memories.length),
+  };
+}

src/services/memory/index.js

@@ -65,7 +65,7 @@ export async function getRelevantMemories(params) {
     return cached;
   }
 
-  const ranked = filterAndRankMemories([...memoryStore.values()], queryParams);
+  const ranked = await filterAndRankMemories([...memoryStore.values()], queryParams);
   await updateMemoryUsage(repository, ranked.map((memory) => memory.id));
   cache.set(cacheKey, ranked);
   return ranked;

src/services/memory/retriever.js

@@ -1,43 +1,170 @@
+const DEFAULT_WEIGHTS = {
+  w_sim: 0.6,
+  w_imp: 0.25,
+  w_rec: 0.15,
+  pinBoost: 0.15,
+};
+
+const DEFAULT_RECENCY_HALF_LIFE_MS = 1000 * 60 * 60 * 24 * 7;
+
+function cosineSimilarity(left = [], right = []) {
+  const size = Math.max(left.length, right.length);
+  let dot = 0;
+  let leftNorm = 0;
+  let rightNorm = 0;
+
+  for (let index = 0; index < size; index += 1) {
+    const leftValue = Number(left[index] ?? 0);
+    const rightValue = Number(right[index] ?? 0);
+    dot += leftValue * rightValue;
+    leftNorm += leftValue ** 2;
+    rightNorm += rightValue ** 2;
+  }
+
+  if (leftNorm === 0 || rightNorm === 0) return 0;
+  return dot / (Math.sqrt(leftNorm) * Math.sqrt(rightNorm));
+}
+
+function computeRecencyDecay(timestamp, now, halfLifeMs) {
+  if (!Number.isFinite(timestamp) || !Number.isFinite(now) || halfLifeMs <= 0) return 0;
+  const age = Math.max(0, now - timestamp);
+  return Math.exp(-(age / halfLifeMs));
+}
+
+function textMatchSimilarity(query, memory) {
+  const normalizedQuery = String(query ?? '').trim().toLowerCase();
+  if (!normalizedQuery) return 0;
+
+  const haystack = `${memory.summary} ${memory.content}`.toLowerCase();
+  if (haystack.includes(normalizedQuery)) return 1;
+
+  const queryTerms = normalizedQuery.split(/\s+/).filter(Boolean);
+  if (queryTerms.length === 0) return 0;
+
+  const matchedTerms = queryTerms.filter((term) => haystack.includes(term)).length;
+  return matchedTerms / queryTerms.length;
+}
+
+/**
+ * @param {import('./schema.js').MemoryRecord} memory
+ * @param {{ tags: string[], timeframe?: { from?: number, to?: number }, pinnedPreference?: 'prefer' | 'only' | 'exclude' }} params
+ */
+function matchesMetadataFilters(memory, params) {
+  const inTimeframe = !params.timeframe
+    || ((params.timeframe.from == null || memory.created_at >= params.timeframe.from)
+      && (params.timeframe.to == null || memory.created_at <= params.timeframe.to));
+
+  if (!inTimeframe) return false;
+
+  if (params.tags.length > 0) {
+    const memoryTags = new Set(memory.tags);
+    const matchesAllTags = params.tags.every((tag) => memoryTags.has(tag));
+    if (!matchesAllTags) return false;
+  }
+
+  if (params.pinnedPreference === 'only' && !memory.pinned) return false;
+  if (params.pinnedPreference === 'exclude' && memory.pinned) return false;
+
+  return true;
+}
+
 /**
  * @param {import('./schema.js').MemoryRecord[]} memories
- * @param {{ agent_id: string, query?: string, tags?: string[], timeframe?: { from?: number, to?: number }, k?: number }} params
- * @returns {import('./schema.js').MemoryRecord[]}
+ * @param {{ query: string, agent_id: string, k: number, vectorAdapter?: { embedText: (text: string) => Promise<number[]> | number[], queryVector: (input: { vector: number[], k?: number, filter?: Record<string, unknown> }) => Promise<Array<{ id: string, score: number }>> | Array<{ id: string, score: number }> }, embedText?: (value: string) => Promise<number[]> | number[] }} params
  */
-export function filterAndRankMemories(memories, params) {
-  const { agent_id, query = '', tags = [], timeframe, k = 10 } = params;
-  const loweredQuery = query.toLowerCase();
+async function buildSimilarityIndex(memories, params) {
+  const similarityById = new Map();
+  const normalizedQuery = params.query.trim();
+
+  if (!normalizedQuery) {
+    return similarityById;
+  }
+
+  if (params.vectorAdapter && typeof params.vectorAdapter.queryVector === 'function') {
+    const queryVector = typeof params.embedText === 'function'
+      ? await params.embedText(normalizedQuery)
+      : await params.vectorAdapter.embedText(normalizedQuery);
+
+    const vectorHits = await params.vectorAdapter.queryVector({
+      vector: queryVector,
+      k: Math.max(params.k * 4, memories.length),
+      filter: { agent_id: params.agent_id },
+    });
+
+    for (const hit of vectorHits) {
+      if (typeof hit.id === 'string' && Number.isFinite(hit.score)) {
+        similarityById.set(hit.id, Math.max(0, Number(hit.score)));
+      }
+    }
+
+    if (similarityById.size > 0) {
+      return similarityById;
+    }
+  }
+
+  for (const memory of memories) {
+    const lexicalScore = textMatchSimilarity(normalizedQuery, memory);
+    if (lexicalScore > 0) {
+      similarityById.set(memory.id, lexicalScore);
+    }
+  }
+
+  return similarityById;
+}
+
+/**
+ * @param {import('./schema.js').MemoryRecord[]} memories
+ * @param {{ agent_id: string, query?: string, tags?: string[], timeframe?: { from?: number, to?: number }, k?: number, pinnedPreference?: 'prefer' | 'only' | 'exclude', weights?: { w_sim?: number, w_imp?: number, w_rec?: number, pinBoost?: number }, now?: number, recencyHalfLifeMs?: number, vectorAdapter?: { embedText: (text: string) => Promise<number[]> | number[], queryVector: (input: { vector: number[], k?: number, filter?: Record<string, unknown> }) => Promise<Array<{ id: string, score: number }>> | Array<{ id: string, score: number }> }, embedText?: (value: string) => Promise<number[]> | number[] }} params
+ * @returns {Promise<import('./schema.js').MemoryRecord[]>}
+ */
+export async function filterAndRankMemories(memories, params) {
+  const {
+    agent_id,
+    query = '',
+    tags = [],
+    timeframe,
+    k = 10,
+    pinnedPreference = 'prefer',
+    weights = {},
+    now = Date.now(),
+    recencyHalfLifeMs = DEFAULT_RECENCY_HALF_LIFE_MS,
+    vectorAdapter,
+    embedText,
+  } = params;
+
+  const mergedWeights = {
+    ...DEFAULT_WEIGHTS,
+    ...weights,
+  };
 
   const filtered = memories.filter((memory) => {
     if (memory.agent_id !== agent_id) return false;
+    return matchesMetadataFilters(memory, { tags, timeframe, pinnedPreference });
+  });
 
-    const inTimeframe = !timeframe
-      || ((timeframe.from == null || memory.created_at >= timeframe.from)
-        && (timeframe.to == null || memory.created_at <= timeframe.to));
-
-    if (!inTimeframe) return false;
-
-    if (tags.length > 0) {
-      const memoryTags = new Set(memory.tags);
-      const matchesAllTags = tags.every((tag) => memoryTags.has(tag));
-      if (!matchesAllTags) return false;
-    }
-
-    return true;
+  const similarityById = await buildSimilarityIndex(filtered, {
+    query,
+    agent_id,
+    k,
+    vectorAdapter,
+    embedText,
   });
 
   return filtered
     .map((memory) => {
-      const haystack = `${memory.summary} ${memory.content}`.toLowerCase();
-      const queryScore = loweredQuery && haystack.includes(loweredQuery) ? 1 : 0;
-      const pinBoost = memory.pinned ? 1 : 0;
-      const recencyScore = Number.isFinite(memory.last_seen) ? memory.last_seen : memory.created_at;
+      const memoryTimestamp = Number.isFinite(memory.last_seen) ? memory.last_seen : memory.created_at;
+      const recencyDecay = computeRecencyDecay(memoryTimestamp, now, recencyHalfLifeMs);
+      const similarity = similarityById.get(memory.id) ?? 0;
+      const pinBoost = memory.pinned && pinnedPreference !== 'exclude' ? mergedWeights.pinBoost : 0;
 
-      return {
-        memory,
-        score: (queryScore * 10) + (pinBoost * 5) + recencyScore,
-      };
+      const score = (similarity * mergedWeights.w_sim)
+        + (memory.importance * mergedWeights.w_imp)
+        + (recencyDecay * mergedWeights.w_rec)
+        + pinBoost;
+
+      return { memory, score };
     })
-    .sort((a, b) => b.score - a.score)
+    .sort((left, right) => right.score - left.score)
     .slice(0, k)
     .map((entry) => entry.memory);
 }
@@ -50,3 +177,5 @@ export function filterAndRankMemories(memories, params) {
 export async function updateMemoryUsage(repository, ids, lastSeen = Date.now()) {
   await Promise.all(ids.map((id) => repository.updateMemoryUsage(id, lastSeen)));
 }
+
+export { cosineSimilarity };

test/memory-retriever.test.mjs

@@ -0,0 +1,129 @@
+import test from 'node:test';
+import assert from 'node:assert/strict';
+
+import { filterAndRankMemories } from '../src/services/memory/retriever.js';
+import { formatMemoriesForPrompt } from '../src/services/memory/formatter.js';
+
+function createMemory(overrides = {}) {
+  return {
+    schema_version: 1,
+    id: overrides.id ?? crypto.randomUUID(),
+    agent_id: overrides.agent_id ?? 'agent-1',
+    session_id: overrides.session_id ?? 's1',
+    type: overrides.type ?? 'event',
+    content: overrides.content ?? 'default content',
+    summary: overrides.summary ?? 'default summary',
+    tags: overrides.tags ?? [],
+    importance: overrides.importance ?? 0.5,
+    embedding_id: overrides.embedding_id ?? null,
+    created_at: overrides.created_at ?? 1000,
+    last_seen: overrides.last_seen ?? 1000,
+    source: overrides.source ?? 'ingest',
+    ttl_days: overrides.ttl_days ?? null,
+    merged_into: overrides.merged_into ?? null,
+    pinned: overrides.pinned ?? false,
+  };
+}
+
+test('filterAndRankMemories applies metadata filters and composite score', async () => {
+  const now = 10_000;
+  const memories = [
+    createMemory({
+      id: 'm1',
+      content: 'billing incident root cause',
+      summary: 'billing incident',
+      importance: 0.9,
+      tags: ['incident', 'billing'],
+      created_at: 9000,
+      last_seen: 9500,
+      pinned: true,
+    }),
+    createMemory({
+      id: 'm2',
+      content: 'billing note',
+      summary: 'routine billing note',
+      importance: 0.2,
+      tags: ['billing'],
+      created_at: 8000,
+      last_seen: 8000,
+      pinned: false,
+    }),
+    createMemory({
+      id: 'm3',
+      content: 'other agent record',
+      summary: 'not for this agent',
+      agent_id: 'agent-2',
+      tags: ['incident', 'billing'],
+    }),
+  ];
+
+  const vectorAdapter = {
+    async embedText() {
+      return [1, 0];
+    },
+    async queryVector() {
+      return [
+        { id: 'm1', score: 0.95 },
+        { id: 'm2', score: 0.8 },
+      ];
+    },
+  };
+
+  const ranked = await filterAndRankMemories(memories, {
+    agent_id: 'agent-1',
+    query: 'billing incident',
+    tags: ['billing'],
+    timeframe: { from: 7000, to: 10_000 },
+    pinnedPreference: 'prefer',
+    k: 5,
+    now,
+    vectorAdapter,
+  });
+
+  assert.deepEqual(ranked.map((item) => item.id), ['m1', 'm2']);
+
+  const pinnedOnly = await filterAndRankMemories(memories, {
+    agent_id: 'agent-1',
+    query: 'billing',
+    pinnedPreference: 'only',
+    k: 5,
+    now,
+    vectorAdapter,
+  });
+
+  assert.deepEqual(pinnedOnly.map((item) => item.id), ['m1']);
+});
+
+test('formatMemoriesForPrompt returns prompt blocks and enforces token budget trimming', () => {
+  const memories = [
+    createMemory({
+      id: 'm1',
+      created_at: 1_700_000_000_000,
+      importance: 0.95,
+      summary: 'A'.repeat(200),
+      content: 'B'.repeat(500),
+      tags: ['incident', 'urgent'],
+    }),
+    createMemory({
+      id: 'm2',
+      created_at: 1_700_000_100_000,
+      importance: 0.7,
+      summary: 'second',
+      content: 'short context',
+      tags: ['ops'],
+    }),
+  ];
+
+  const result = formatMemoriesForPrompt(memories, {
+    tokenBudget: 70,
+    k: 2,
+  });
+
+  assert.ok(result.blocks.length >= 1);
+  assert.ok(result.blocks.length <= 2);
+  assert.ok(result.usedTokens <= 70);
+  assert.ok(result.text.includes('[m1,'));
+  if (result.blocks.length === 1) {
+    assert.equal(result.truncated, true);
+  }
+});