diff --git a/scripts/benchmark-pruner.mjs b/scripts/benchmark-pruner.mjs new file mode 100644 index 0000000..dae4f89 --- /dev/null +++ b/scripts/benchmark-pruner.mjs @@ -0,0 +1,135 @@ +import { createLifecyclePlan } from '../src/services/memory/pruner.js'; +import crypto from 'node:crypto'; +import { performance } from 'node:perf_hooks'; + +const ITERATIONS = 3; +const GROUPS = 100; +const GROUP_SIZE = 5; +const LATENCY_MS = 10; + +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 ?? Date.now(), + last_seen: overrides.last_seen ?? Date.now(), + source: overrides.source ?? 'ingest', + ttl_days: overrides.ttl_days ?? null, + merged_into: overrides.merged_into ?? null, + pinned: overrides.pinned ?? false, + }; +} + +function generateMemories() { + const memories = []; + const now = Date.now(); + + for (let i = 0; i < GROUPS; i++) { + const groupId = `group-${i}`; + for (let j = 0; j < GROUP_SIZE; j++) { + memories.push(createMemory({ + id: `${groupId}-${j}`, + content: `Memory ${j} for group ${groupId}`, + summary: `Summary ${j} for group ${groupId}`, + tags: ['benchmark', groupId], + created_at: now - (j * 100), // Close in time + last_seen: now, + importance: 0.5, + })); + } + } + + // Add some random noise memories that won't group + for (let i = 0; i < 50; i++) { + memories.push(createMemory({ + id: `noise-${i}`, + content: `Noise memory ${i}`, + tags: [`noise-${i}`], // Unique tags to prevent grouping + created_at: now - (i * 100000), // Spread out + })); + } + + return memories; +} + +const delay = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); + +async function runBenchmark() { + console.log(`Setting up benchmark: ${GROUPS} groups of ${GROUP_SIZE} memories + noise.`); + console.log(`Mock latency: ${LATENCY_MS}ms per group condensation.`); + + const memories = generateMemories(); + + // Prepare embeddings + const groupEmbeddings = new Map(); + for (let i = 0; i < GROUPS; i++) { + const vec = new Array(GROUPS).fill(0); + vec[i] = 1; + groupEmbeddings.set(`group-${i}`, vec); + } + + const getEmbeddingWrapper = (memory) => { + const groupIdTag = memory.tags.find(t => t.startsWith('group-')); + if (groupIdTag && groupEmbeddings.has(groupIdTag)) { + return groupEmbeddings.get(groupIdTag); + } + // Return a unique random vector for noise to prevent grouping + // But since they have unique tags, getEmbedding won't even be called for similarity check + // if tag overlap check fails first? + // groupNearDuplicates logic: + // ... if (!withinWindow || !hasTagOverlap(base.tags, probe.tags)) continue; + // So if tags don't overlap, it continues. + // So embedding doesn't matter for noise if tags are unique. + return [Math.random()]; + }; + + const options = { + retentionMs: 30 * 24 * 60 * 60 * 1000, // 30 days + duplicateWindowMs: 1000 * 60 * 60, // 1 hour + similarityThreshold: 0.99, + getEmbedding: getEmbeddingWrapper, + generateSummary: async (prompt) => { + await delay(LATENCY_MS); + return JSON.stringify({ + summary: 'Condensed summary based on prompt.', + importance: 0.8, + }); + }, + // Mock templateDir to avoid file reading issues if paths are tricky, + // but loadMemoryPromptTemplates reads from file system. + // It should work if run from root. + }; + + console.log(`Starting ${ITERATIONS} iterations...`); + + const times = []; + + for (let i = 0; i < ITERATIONS; i++) { + const start = performance.now(); + const plan = await createLifecyclePlan(memories, options); + const duration = performance.now() - start; + + // Verify grouping happened + const condensedCount = plan.condensedGroups.length; + // We expect exactly GROUPS condensed groups. + if (condensedCount !== GROUPS) { + console.warn(`Warning: Expected ${GROUPS} condensed groups, but got ${condensedCount}. Check grouping logic.`); + } + + times.push(duration); + console.log(`Iteration ${i + 1}: ${duration.toFixed(2)}ms`); + } + + const avg = times.reduce((a, b) => a + b, 0) / times.length; + console.log(`Average execution time: ${avg.toFixed(2)}ms`); +} + +runBenchmark().catch(console.error); diff --git a/src/services/memory/pruner.js b/src/services/memory/pruner.js index bf996d6..1c9aeae 100644 --- a/src/services/memory/pruner.js +++ b/src/services/memory/pruner.js @@ -189,12 +189,13 @@ export async function createLifecyclePlan(memories, options) { }); const templates = loadMemoryPromptTemplates({ templateDir: options.templateDir }); - const condensedGroups = []; + const condensedGroups = await Promise.all( + groups.map((group) => condenseGroup(group, options.generateSummary, templates.condense)) + ); + const mergedIntoMap = new Map(); - for (const group of groups) { - const condensed = await condenseGroup(group, options.generateSummary, templates.condense); - condensedGroups.push(condensed); + for (const condensed of condensedGroups) { for (const mergedId of condensed.mergedIds) { mergedIntoMap.set(mergedId, condensed.merged.id); }