diff --git a/src/services/memory/formatter.js b/src/services/memory/formatter.js new file mode 100644 index 0000000..9ebbd7e --- /dev/null +++ b/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), + }; +} diff --git a/src/services/memory/index.js b/src/services/memory/index.js index 95a93c1..8752d9d 100644 --- a/src/services/memory/index.js +++ b/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; diff --git a/src/services/memory/retriever.js b/src/services/memory/retriever.js index 5bb9031..8b3d76f 100644 --- a/src/services/memory/retriever.js +++ b/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[], queryVector: (input: { vector: number[], k?: number, filter?: Record }) => Promise> | Array<{ id: string, score: number }> }, embedText?: (value: string) => Promise | 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[], queryVector: (input: { vector: number[], k?: number, filter?: Record }) => Promise> | Array<{ id: string, score: number }> }, embedText?: (value: string) => Promise | number[] }} params + * @returns {Promise} + */ +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 }; diff --git a/test/memory-retriever.test.mjs b/test/memory-retriever.test.mjs new file mode 100644 index 0000000..8c14327 --- /dev/null +++ b/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); + } +});