diff --git a/src/services/memory/formatter.js b/src/services/memory/formatter.js index 9ebbd7e..8b04cc1 100644 --- a/src/services/memory/formatter.js +++ b/src/services/memory/formatter.js @@ -1,6 +1,7 @@ const DEFAULT_TOKEN_BUDGET = 1000; const DEFAULT_MAX_SUMMARY_CHARS = 240; const DEFAULT_MAX_CONTEXT_CHARS = 320; +const MEMORY_CONTEXT_HEADER = '[CONTEXT — Relevant memories]'; function clampString(value, maxChars) { if (typeof value !== 'string') return ''; @@ -104,3 +105,114 @@ export function formatMemoriesForPrompt(memories, options = {}) { truncated: blocks.length < Math.min(k, memories.length), }; } + +function toContextString(agentContext) { + if (typeof agentContext === 'string') { + return agentContext.trim(); + } + + if (agentContext == null) { + return ''; + } + + if (Array.isArray(agentContext)) { + return agentContext.map((value) => String(value)).join('\n').trim(); + } + + if (typeof agentContext === 'object') { + try { + return JSON.stringify(agentContext); + } catch { + return String(agentContext).trim(); + } + } + + return String(agentContext).trim(); +} + +/** + * Supports safe rollout and canarying with one env var: + * - true/1/on/all => enabled for all agents + * - false/0/off/no/empty => disabled + * - comma-separated agent ids => enabled only for matching ids + * + * @param {string} agentId + * @param {Record} [env] + */ +export function isMemoryRetrievalEnabled(agentId, env = process.env) { + const raw = String(env.TORCH_MEMORY_RETRIEVAL_ENABLED ?? '').trim(); + if (!raw) return false; + + const normalized = raw.toLowerCase(); + if (['0', 'false', 'off', 'no', 'disabled'].includes(normalized)) return false; + if (['1', 'true', 'on', 'yes', 'enabled', 'all'].includes(normalized)) return true; + + const allowList = raw + .split(',') + .map((value) => value.trim()) + .filter(Boolean); + + if (allowList.length === 0) return false; + return allowList.includes(agentId); +} + +/** + * Builds a final prompt with retrieved memory context inserted ahead of the base prompt. + * + * @param {{ + * basePrompt: string, + * userRequest: string, + * agentContext?: unknown, + * agent_id: string, + * memoryService: { getRelevantMemories: (params: { agent_id: string, query: string, k?: number }) => Promise | import('./schema.js').MemoryRecord[] }, + * retrievalK?: number, + * tokenBudget?: number, + * env?: Record + * }} params + */ +export async function renderPromptWithMemoryContext(params) { + const { + basePrompt, + userRequest, + agentContext, + agent_id, + memoryService, + retrievalK = 5, + tokenBudget = 320, + env = process.env, + } = params; + + if (!isMemoryRetrievalEnabled(agent_id, env)) { + return String(basePrompt ?? ''); + } + + const requestText = String(userRequest ?? '').trim(); + const contextText = toContextString(agentContext); + const query = [requestText, contextText].filter(Boolean).join('\n\n'); + + if (!query || !memoryService?.getRelevantMemories) { + return String(basePrompt ?? ''); + } + + const memories = await memoryService.getRelevantMemories({ + agent_id, + query, + k: retrievalK, + }); + + if (memoryService?.updateMemoryUsage) { + const now = Date.now(); + await Promise.all((memories ?? []).map((memory) => memoryService.updateMemoryUsage(memory.id, now))); + } + + const rendered = formatMemoriesForPrompt(memories ?? [], { + k: retrievalK, + tokenBudget, + }); + + if (!rendered.text) { + return String(basePrompt ?? ''); + } + + return `${MEMORY_CONTEXT_HEADER}\n${rendered.text}\n\n${String(basePrompt ?? '')}`; +} diff --git a/test/memory-retriever.test.mjs b/test/memory-retriever.test.mjs index 8c14327..9fee8e6 100644 --- a/test/memory-retriever.test.mjs +++ b/test/memory-retriever.test.mjs @@ -2,7 +2,11 @@ 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'; +import { + formatMemoriesForPrompt, + isMemoryRetrievalEnabled, + renderPromptWithMemoryContext, +} from '../src/services/memory/formatter.js'; function createMemory(overrides = {}) { return { @@ -127,3 +131,68 @@ test('formatMemoriesForPrompt returns prompt blocks and enforces token budget tr assert.equal(result.truncated, true); } }); + +test('isMemoryRetrievalEnabled supports global toggle and agent canary list', () => { + assert.equal(isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'true' }), true); + assert.equal(isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'false' }), false); + assert.equal( + isMemoryRetrievalEnabled('agent-a', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' }), + true, + ); + assert.equal( + isMemoryRetrievalEnabled('agent-c', { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' }), + false, + ); +}); + +test('renderPromptWithMemoryContext injects bounded memory block before base prompt', async () => { + const memories = [ + createMemory({ + id: 'm-ctx', + summary: 'Known issue in deploy pipeline', + content: 'Retry build after rotating temporary token', + created_at: 1_700_000_000_000, + importance: 0.91, + tags: ['deploy', 'incident'], + }), + ]; + + const calls = []; + const usageUpdates = []; + const memoryService = { + async getRelevantMemories(params) { + calls.push(params); + return memories; + }, + async updateMemoryUsage(id, lastSeen) { + usageUpdates.push({ id, lastSeen }); + }, + }; + + const output = await renderPromptWithMemoryContext({ + basePrompt: 'Base final prompt body', + userRequest: 'Please debug deployment failure', + agentContext: { branch: 'main', tool: 'torch' }, + agent_id: 'agent-a', + memoryService, + retrievalK: 1, + tokenBudget: 60, + env: { TORCH_MEMORY_RETRIEVAL_ENABLED: 'agent-a,agent-b' }, + }); + + assert.equal(calls.length, 1); + assert.equal(calls[0].agent_id, 'agent-a'); + assert.equal(calls[0].k, 1); + assert.match(calls[0].query, /Please debug deployment failure/); + assert.match(calls[0].query, /"branch":"main"/); + + const headerIndex = output.indexOf('[CONTEXT — Relevant memories]'); + const baseIndex = output.indexOf('Base final prompt body'); + assert.ok(headerIndex >= 0); + assert.ok(baseIndex > headerIndex); + assert.ok(output.includes('[m-ctx,')); + + assert.equal(usageUpdates.length, 1); + assert.equal(usageUpdates[0].id, 'm-ctx'); + assert.equal(typeof usageUpdates[0].lastSeen, 'number'); +});