diff --git a/packages/agent-memory/src/AgentMemory.ts b/packages/agent-memory/src/AgentMemory.ts index 560ddbfd..c451ddf0 100644 --- a/packages/agent-memory/src/AgentMemory.ts +++ b/packages/agent-memory/src/AgentMemory.ts @@ -1 +1,86 @@ -export class AgentMemory {} +import { AgentCache, type AgentCacheOptions } from '@betterdb/agent-cache'; +import { + MemoryStore, + type MemoryDiscoveryConfig, + type MemoryConfigRefreshConfig, +} from './MemoryStore'; +import type { RecallWeights } from './compositeScore'; +import type { EmbedFn, MemoryStoreClient } from './types'; + +const DEFAULT_NAME = 'betterdb_ac'; + +export interface AgentMemoryConfig { + defaultThreshold?: number; + recall?: { + weights?: RecallWeights; + halfLifeSeconds?: number; + }; + maxItemsPerScope?: number; + discovery?: boolean | MemoryDiscoveryConfig; + configRefresh?: boolean | MemoryConfigRefreshConfig; +} + +export interface AgentMemoryOptions extends AgentCacheOptions { + embedFn: EmbedFn; + memory?: AgentMemoryConfig; +} + +export class AgentMemory { + readonly llm: AgentCache['llm']; + readonly tool: AgentCache['tool']; + readonly session: AgentCache['session']; + readonly memory: MemoryStore; + private readonly cache: AgentCache; + + constructor(options: AgentMemoryOptions) { + if (typeof options.embedFn !== 'function') { + throw new Error('AgentMemory requires an embedFn to back the memory tier'); + } + + // Resolve the name once and hand the same value to both tiers so their key + // prefixes, discovery markers, and stats keys can never drift apart. + const name = options.name ?? DEFAULT_NAME; + this.cache = new AgentCache({ ...options, name }); + this.llm = this.cache.llm; + this.tool = this.cache.tool; + this.session = this.cache.session; + + const memory = options.memory ?? {}; + this.memory = new MemoryStore({ + // AgentCacheOptions.client doesn't surface the `.call` method MemoryStore + // needs; a real ioredis/iovalkey client has it, so we assert the contract + // here. A method-only client/mock would compile but fail at runtime. + client: options.client as unknown as MemoryStoreClient, + name, + embedFn: options.embedFn, + defaultThreshold: memory.defaultThreshold, + weights: memory.recall?.weights, + halfLifeSeconds: memory.recall?.halfLifeSeconds, + maxItemsPerScope: memory.maxItemsPerScope, + // The facade is the batteries-included product: discover the memory tier + // alongside the cache tiers by default, unless explicitly disabled. + discovery: memory.discovery ?? true, + configRefresh: memory.configRefresh, + telemetry: options.telemetry?.registry ? { registry: options.telemetry.registry } : undefined, + }); + } + + async initialize(): Promise { + // Surface a discovery name-collision from either tier: awaiting + // ensureDiscoveryReady() is AgentCache's documented strict collision check, + // and the memory tier already propagates, so the cache side isn't swallowed. + await Promise.all([ + this.cache.ensureDiscoveryReady(), + this.memory.ensureDiscoveryReady(), + ]); + } + + async close(): Promise { + // Tear down both tiers even if one fails, so timers and heartbeats can't leak. + try { + await this.memory.close(); + } finally { + await this.cache.shutdown(); + } + } +} diff --git a/packages/agent-memory/src/MemoryStore.ts b/packages/agent-memory/src/MemoryStore.ts index 02d77961..2edb18d3 100644 --- a/packages/agent-memory/src/MemoryStore.ts +++ b/packages/agent-memory/src/MemoryStore.ts @@ -243,6 +243,12 @@ export class MemoryStore { return discovery; } + async ensureDiscoveryReady(): Promise { + if (this.discoveryReady) { + await this.discoveryReady.catch(() => undefined); + } + } + async close(): Promise { if (this.configRefreshHandle) { clearInterval(this.configRefreshHandle); diff --git a/packages/agent-memory/src/__tests__/AgentMemory.test.ts b/packages/agent-memory/src/__tests__/AgentMemory.test.ts new file mode 100644 index 00000000..c9736ce4 --- /dev/null +++ b/packages/agent-memory/src/__tests__/AgentMemory.test.ts @@ -0,0 +1,185 @@ +import { describe, it, expect, vi } from 'vitest'; +import { Registry } from 'prom-client'; +import { AgentMemory, type AgentMemoryOptions } from '../AgentMemory'; +import { MemoryStore } from '../MemoryStore'; +import { fakeEmbed } from './helpers/fakeEmbed'; + +function fakeValkey() { + const ok = vi.fn(async () => 'OK'); + const nul = vi.fn(async () => null); + return { + call: vi.fn(async () => 'OK'), + get: nul, + set: ok, + del: ok, + hget: nul, + hset: ok, + hgetall: vi.fn(async () => ({})), + hincrby: ok, + expire: ok, + exists: vi.fn(async () => 0), + scan: vi.fn(async () => ['0', []]), + }; +} + +type FakeClient = ReturnType; + +function makeOptions(overrides: Partial = {}): AgentMemoryOptions { + return { + client: fakeValkey() as unknown as AgentMemoryOptions['client'], + embedFn: fakeEmbed(8), + discovery: { enabled: false }, + configRefresh: { enabled: false }, + analytics: { disabled: true }, + ...overrides, + } as AgentMemoryOptions; +} + +describe('AgentMemory facade', () => { + it('exposes the three short-term tiers plus the memory tier', async () => { + const mem = new AgentMemory(makeOptions()); + + expect(mem.llm).toBeDefined(); + expect(mem.tool).toBeDefined(); + expect(mem.session).toBeDefined(); + expect(mem.memory).toBeInstanceOf(MemoryStore); + + await mem.close(); + }); + + it('throws a clear error when constructed without an embedFn', () => { + const options = { ...makeOptions(), embedFn: undefined } as unknown as AgentMemoryOptions; + expect(() => new AgentMemory(options)).toThrow(/embedFn/i); + }); + + it('wires the memory tier to the shared client and default prefix', async () => { + const client = fakeValkey(); + const mem = new AgentMemory( + makeOptions({ client: client as unknown as AgentMemoryOptions['client'] }), + ); + + const id = await mem.memory.remember('hello'); + + expect(typeof id).toBe('string'); + const hset = (client.call.mock.calls as unknown[][]).find( + (c) => c[0] === 'HSET' && typeof c[1] === 'string' && c[1].startsWith('betterdb_ac:mem:'), + ); + expect(hset).toBeDefined(); + + await mem.close(); + }); + + it('shares the configured name as the memory key prefix', async () => { + const client = fakeValkey(); + const mem = new AgentMemory( + makeOptions({ client: client as unknown as AgentMemoryOptions['client'], name: 'myapp' }), + ); + + await mem.memory.remember('hello'); + + const hset = (client.call.mock.calls as unknown[][]).find( + (c) => c[0] === 'HSET' && typeof c[1] === 'string' && c[1].startsWith('myapp:mem:'), + ); + expect(hset).toBeDefined(); + + await mem.close(); + }); + + it('maps the memory sub-config onto the MemoryStore', async () => { + const mem = new AgentMemory( + makeOptions({ + memory: { + defaultThreshold: 0.4, + recall: { + weights: { similarity: 0.5, recency: 0.3, importance: 0.2 }, + halfLifeSeconds: 3600, + }, + maxItemsPerScope: 100, + }, + }), + ); + + expect(mem.memory.currentConfig()).toEqual({ + threshold: 0.4, + weights: { similarity: 0.5, recency: 0.3, importance: 0.2 }, + halfLifeSeconds: 3600, + maxItemsPerScope: 100, + }); + + await mem.close(); + }); + + it('initialize() resolves and close() tears down both tiers', async () => { + const mem = new AgentMemory(makeOptions()); + const memoryClose = vi.spyOn(mem.memory, 'close'); + + await expect(mem.initialize()).resolves.toBeUndefined(); + await mem.close(); + + expect(memoryClose).toHaveBeenCalled(); + }); + + it('initialize() surfaces a cache discovery collision instead of swallowing it', async () => { + const mem = new AgentMemory(makeOptions()); + const cache = (mem as unknown as { cache: { ensureDiscoveryReady: () => Promise } }).cache; + vi.spyOn(cache, 'ensureDiscoveryReady').mockRejectedValue(new Error('cache name collision')); + + await expect(mem.initialize()).rejects.toThrow(/collision/i); + + await mem.close(); + }); + + it('registers a memory discovery marker by default', async () => { + const client = fakeValkey(); + const mem = new AgentMemory( + makeOptions({ client: client as unknown as AgentMemoryOptions['client'] }), + ); + + await mem.initialize(); + + const marker = (client.call.mock.calls as unknown[][]).find( + (c) => c[0] === 'HSET' && c[1] === '__betterdb:caches', + ); + expect(marker).toBeDefined(); + expect(JSON.parse(marker?.[3] as string).type).toBe('agent_memory'); + // The memory marker registers under a distinct `{name}:mem` field so it + // can't clobber an agent_cache marker sharing the same name. + expect(marker?.[2]).toBe('betterdb_ac:mem'); + + await mem.close(); + }); + + it('allows disabling memory discovery', async () => { + const client = fakeValkey(); + const mem = new AgentMemory( + makeOptions({ + client: client as unknown as AgentMemoryOptions['client'], + memory: { discovery: false }, + }), + ); + + await mem.initialize(); + await mem.close(); + + const marker = (client.call.mock.calls as unknown[][]).find( + (c) => c[0] === 'HSET' && c[1] === '__betterdb:caches', + ); + expect(marker).toBeUndefined(); + }); + + it('shares one prom registry across the cache and memory tiers', async () => { + const registry = new Registry(); + const mem = new AgentMemory(makeOptions({ telemetry: { registry } })); + + await mem.memory.remember('x'); + + const text = await registry.metrics(); + expect(text).toMatch(/agent_memory_embedding_calls_total/); + expect(text).toMatch(/agent_cache_/); + + await mem.close(); + }); +}); + +// Touch the FakeClient type so it is exercised by the suite. +export type { FakeClient }; diff --git a/packages/agent-memory/src/index.ts b/packages/agent-memory/src/index.ts index 87552912..7ac3f056 100644 --- a/packages/agent-memory/src/index.ts +++ b/packages/agent-memory/src/index.ts @@ -11,6 +11,7 @@ export type { MemoryDiscoveryDeps, MemoryMarker } from './discovery'; export { createMemoryTelemetry, DEFAULT_METRICS_PREFIX, DEFAULT_TRACER_NAME } from './telemetry'; export type { MemoryTelemetry, MemoryTelemetryOptions, MemoryMetrics } from './telemetry'; export { AgentMemory } from './AgentMemory'; +export type { AgentMemoryOptions, AgentMemoryConfig } from './AgentMemory'; export type { EmbedFn, MemoryStoreClient,