diff --git a/strands-ts/package.json b/strands-ts/package.json index b2f49211b9..24c3e7db9a 100644 --- a/strands-ts/package.json +++ b/strands-ts/package.json @@ -68,6 +68,10 @@ "types": "./dist/src/session/s3-storage.d.ts", "default": "./dist/src/session/s3-storage.js" }, + "./memory/stores/bedrock-knowledge-base": { + "types": "./dist/src/memory/stores/bedrock-knowledge-base-store.d.ts", + "default": "./dist/src/memory/stores/bedrock-knowledge-base-store.js" + }, "./telemetry": { "types": "./dist/src/telemetry/index.d.ts", "default": "./dist/src/telemetry/index.js" @@ -140,6 +144,8 @@ "@ai-sdk/provider": "^3.0.0", "@anthropic-ai/sdk": "^0.92.0", "@aws-sdk/client-bedrock": "^3.943.0", + "@aws-sdk/client-bedrock-agent": "^3.943.0", + "@aws-sdk/client-bedrock-agent-runtime": "^3.943.0", "@aws-sdk/client-s3": "^3.943.0", "@aws-sdk/client-secrets-manager": "^3.943.0", "@aws-sdk/client-sts": "^3.996.0", @@ -193,6 +199,8 @@ "@a2a-js/sdk": "^0.3.10", "@ai-sdk/provider": "^3.0.0", "@anthropic-ai/sdk": "^0.92.0", + "@aws-sdk/client-bedrock-agent": "^3.943.0", + "@aws-sdk/client-bedrock-agent-runtime": "^3.943.0", "@aws-sdk/client-s3": "^3.943.0", "@aws/bedrock-token-generator": "^1.1.0", "@google/genai": "^1.40.0", @@ -231,6 +239,12 @@ "@aws-sdk/client-s3": { "optional": true }, + "@aws-sdk/client-bedrock-agent": { + "optional": true + }, + "@aws-sdk/client-bedrock-agent-runtime": { + "optional": true + }, "@google/genai": { "optional": true }, diff --git a/strands-ts/src/agent/agent.ts b/strands-ts/src/agent/agent.ts index 41c674ed7a..862d314bcb 100644 --- a/strands-ts/src/agent/agent.ts +++ b/strands-ts/src/agent/agent.ts @@ -73,6 +73,8 @@ import { ToolCaller } from './tool-caller.js' import type { ToolCallerProxy } from './tool-caller.js' import type { z } from 'zod' +import { MemoryManager } from '../memory/memory-manager.js' +import type { MemoryManagerConfig } from '../memory/index.js' import { SessionManager } from '../session/session-manager.js' import { Tracer } from '../telemetry/tracer.js' import { Meter } from '../telemetry/meter.js' @@ -198,6 +200,12 @@ export type AgentConfig = { * Session manager for saving and restoring agent sessions */ sessionManager?: SessionManager + /** + * Memory manager for cross-session knowledge retrieval and storage. + * Manages one or more knowledge stores and exposes search/store tools. + * Accepts a {@link MemoryManager} instance or a {@link MemoryManagerConfig} object (auto-wrapped). + */ + memoryManager?: MemoryManager | MemoryManagerConfig /** * Custom trace attributes to include in all spans. * These attributes are merged with standard attributes in telemetry spans. @@ -287,6 +295,10 @@ export class Agent implements LocalAgent, InvokableAgent { * The session manager for saving and restoring agent sessions, if configured. */ public readonly sessionManager?: SessionManager | undefined + /** + * The memory manager for cross-session knowledge retrieval and storage, if configured. + */ + public readonly memoryManager?: MemoryManager | undefined private readonly _hooksRegistry: HookRegistryImplementation private readonly _pluginRegistry: PluginRegistry @@ -323,6 +335,12 @@ export class Agent implements LocalAgent, InvokableAgent { this.id = config?.id ?? DEFAULT_AGENT_ID if (config?.description !== undefined) this.description = config.description this.sessionManager = config?.sessionManager + this.memoryManager = + config?.memoryManager instanceof MemoryManager + ? config.memoryManager + : config?.memoryManager + ? new MemoryManager(config.memoryManager) + : undefined if (typeof config?.model === 'string') { this.model = new BedrockModel({ modelId: config.model }) @@ -375,6 +393,7 @@ export class Agent implements LocalAgent, InvokableAgent { this._conversationManager, ...retryStrategies, ...(config?.plugins ?? []), + ...(this.memoryManager ? [this.memoryManager] : []), ...(config?.sessionManager ? [config.sessionManager] : []), new ModelPlugin(this.model), ]) diff --git a/strands-ts/src/index.ts b/strands-ts/src/index.ts index 52316b23d0..1ac3e0c311 100644 --- a/strands-ts/src/index.ts +++ b/strands-ts/src/index.ts @@ -286,6 +286,20 @@ export { export { type McpServerConfig } from './mcp-config.js' export type { ElicitationCallback, ElicitationContext } from './types/elicitation.js' +// Memory management +export { MemoryManager } from './memory/index.js' +export type { + MemoryEntry, + MemoryStore, + MemoryStoreConfig, + SearchOptions, + SearchMemoryOptions, + AddMemoryOptions, + MemoryToolConfig, + MemoryManagerConfig, + InjectionConfig, +} from './memory/index.js' + // Session management export { SessionManager } from './session/session-manager.js' export type { diff --git a/strands-ts/src/memory/__tests__/memory-manager.test.ts b/strands-ts/src/memory/__tests__/memory-manager.test.ts new file mode 100644 index 0000000000..39ea547c64 --- /dev/null +++ b/strands-ts/src/memory/__tests__/memory-manager.test.ts @@ -0,0 +1,657 @@ +import { describe, it, expect, vi } from 'vitest' +import { z } from 'zod' +import { Agent } from '../../agent/agent.js' +import { MemoryManager } from '../memory-manager.js' +import { tool } from '../../tools/tool-factory.js' +import type { MemoryStore, MemoryEntry } from '../types.js' +import type { InvokableTool, Tool } from '../../tools/tool.js' +import { logger } from '../../logging/logger.js' +import { Message, TextBlock, ToolUseBlock, ToolResultBlock } from '../../types/messages.js' +import { BeforeModelCallEvent } from '../../hooks/events.js' +import type { LocalAgent } from '../../types/agent.js' + +function createMockStore( + name: string, + options?: { + entries?: MemoryEntry[] + writable?: boolean + description?: string + maxSearchResults?: number + tools?: Tool[] + } +): MemoryStore { + const store: MemoryStore = { + name, + writable: !!options?.writable, + ...(options?.description && { description: options.description }), + ...(options?.maxSearchResults != null && { maxSearchResults: options.maxSearchResults }), + search: vi.fn().mockResolvedValue(options?.entries ?? []), + } + if (options?.writable) { + store.add = vi.fn().mockResolvedValue(undefined) + } + if (options?.tools) { + store.getTools = vi.fn().mockReturnValue(options.tools) + } + return store +} + +function createNamedTool(name: string): Tool { + return tool({ + name, + description: `test tool ${name}`, + inputSchema: z.object({}), + callback: () => 'ok', + }) +} + +describe('MemoryManager', () => { + describe('constructor', () => { + it('throws when stores array is empty', () => { + expect(() => new MemoryManager({ stores: [] })).toThrow('at least one store is required') + }) + + it('creates instance with valid config', () => { + const mm = new MemoryManager({ stores: [createMockStore('test')] }) + expect(mm.name).toBe('strands:memory-manager') + }) + + it('throws when two stores share a name', () => { + expect(() => new MemoryManager({ stores: [createMockStore('dup'), createMockStore('dup')] })).toThrow( + "duplicate store name 'dup'" + ) + }) + + it('throws when a store is writable but has no add method', () => { + const broken: MemoryStore = { name: 'broken', writable: true, search: vi.fn().mockResolvedValue([]) } + expect(() => new MemoryManager({ stores: [broken] })).toThrow("store 'broken' is writable but has no add method") + }) + + it('throws when addToolConfig is enabled but no stores are writable', () => { + expect( + () => + new MemoryManager({ + stores: [createMockStore('a')], + addToolConfig: true, + }) + ).toThrow('addToolConfig is enabled but no stores are writable') + }) + + it('allows addToolConfig true with a single writable store', () => { + const mm = new MemoryManager({ + stores: [createMockStore('a', { writable: true })], + addToolConfig: true, + }) + expect(mm.getTools().map((t) => t.name)).toContain('add_memory') + }) + + it('allows addToolConfig true with multiple writable stores', () => { + const mm = new MemoryManager({ + stores: [createMockStore('a', { writable: true }), createMockStore('b', { writable: true })], + addToolConfig: true, + }) + expect(mm.getTools().map((t) => t.name)).toContain('add_memory') + }) + }) + + describe('getTools', () => { + it('registers search tool by default', () => { + const mm = new MemoryManager({ stores: [createMockStore('test')] }) + const tools = mm.getTools() + expect(tools).toHaveLength(1) + expect(tools[0]!.name).toBe('search_memory') + }) + + it('registers add tool when addToolConfig is enabled', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + addToolConfig: true, + }) + const tools = mm.getTools() + expect(tools.map((t) => t.name)).toStrictEqual(['search_memory', 'add_memory']) + }) + + it('does not register add tool by default', () => { + const mm = new MemoryManager({ stores: [createMockStore('test', { writable: true })] }) + const tools = mm.getTools() + expect(tools.map((t) => t.name)).toStrictEqual(['search_memory']) + }) + + it('returns empty array when searchToolConfig is false and addToolConfig is false', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + searchToolConfig: false, + addToolConfig: false, + }) + expect(mm.getTools()).toStrictEqual([]) + }) + + it('uses custom tool names from MemoryToolConfig', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + searchToolConfig: { name: 'recall' }, + addToolConfig: { name: 'remember' }, + }) + const tools = mm.getTools() + expect(tools.map((t) => t.name)).toStrictEqual(['recall', 'remember']) + }) + + it('includes store descriptions in search tool description', () => { + const store = createMockStore('personal', { description: 'User preferences' }) + const mm = new MemoryManager({ stores: [store] }) + const tools = mm.getTools() + expect(tools[0]!.description).toContain('personal: User preferences') + expect(tools[0]!.description).toContain('target one or more memory stores by name') + }) + + it('includes store descriptions in add tool description', () => { + const store = createMockStore('notes', { writable: true, description: 'Personal notes' }) + const mm = new MemoryManager({ stores: [store], addToolConfig: true }) + const tools = mm.getTools() + const addTool = tools.find((t) => t.name === 'add_memory')! + expect(addTool.description).toContain('notes: Personal notes') + expect(addTool.description).toContain('target a specific store by name') + }) + + it('aggregates tools provided by stores via getTools', () => { + const store = createMockStore('kb', { tools: [createNamedTool('kb_query')] }) + const mm = new MemoryManager({ stores: [store] }) + + expect(mm.getTools().map((t) => t.name)).toStrictEqual(['search_memory', 'kb_query']) + }) + + it('aggregates store tools across multiple stores alongside the manager tools', () => { + const a = createMockStore('a', { writable: true, tools: [createNamedTool('a_tool')] }) + const b = createMockStore('b', { tools: [createNamedTool('b_tool')] }) + const mm = new MemoryManager({ stores: [a, b], addToolConfig: true }) + + expect(mm.getTools().map((t) => t.name)).toStrictEqual(['search_memory', 'add_memory', 'a_tool', 'b_tool']) + }) + + it('includes store tools even when the manager registers no tools of its own', () => { + const store = createMockStore('kb', { tools: [createNamedTool('kb_query')] }) + const mm = new MemoryManager({ stores: [store], searchToolConfig: false }) + + expect(mm.getTools().map((t) => t.name)).toStrictEqual(['kb_query']) + }) + }) + + describe('search', () => { + it('queries all stores and concatenates results', async () => { + const store1 = createMockStore('a', { entries: [{ content: 'fact one' }] }) + const store2 = createMockStore('b', { entries: [{ content: 'fact two' }] }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + const results = await mm.search('query') + expect(results).toStrictEqual([{ content: 'fact one' }, { content: 'fact two' }]) + }) + + it('passes maxSearchResults to each store', async () => { + const store = createMockStore('a', { maxSearchResults: 5 }) + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query') + expect(store.search).toHaveBeenCalledWith('query', { maxSearchResults: 5 }) + }) + + it('overrides per-store maxSearchResults with options.maxSearchResults', async () => { + const store = createMockStore('a', { maxSearchResults: 5 }) + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query', { maxSearchResults: 2 }) + expect(store.search).toHaveBeenCalledWith('query', { maxSearchResults: 2 }) + }) + + it('defaults to 3 results when no maxSearchResults configured', async () => { + const store = createMockStore('a') + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query') + expect(store.search).toHaveBeenCalledWith('query', { maxSearchResults: 3 }) + }) + + it('filters to named stores when options.stores is provided', async () => { + const store1 = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const store2 = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + const results = await mm.search('query', { stores: ['personal'] }) + expect(results).toStrictEqual([{ content: 'personal fact' }]) + expect(store2.search).not.toHaveBeenCalled() + }) + + it('gracefully handles store failures', async () => { + const store1: MemoryStore = { + name: 'failing', + writable: false, + search: vi.fn().mockRejectedValue(new Error('network error')), + } + const store2 = createMockStore('ok', { entries: [{ content: 'fact' }] }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + const results = await mm.search('query') + expect(results).toStrictEqual([{ content: 'fact' }]) + }) + + it('searches all stores when stores option is omitted', async () => { + const store1 = createMockStore('a', { entries: [{ content: 'fact one' }] }) + const store2 = createMockStore('b', { entries: [{ content: 'fact two' }] }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + const results = await mm.search('query') + expect(results).toStrictEqual([{ content: 'fact one' }, { content: 'fact two' }]) + }) + + it('searches no stores when stores option is an empty array', async () => { + const store1 = createMockStore('a', { entries: [{ content: 'fact one' }] }) + const store2 = createMockStore('b', { entries: [{ content: 'fact two' }] }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + const results = await mm.search('query', { stores: [] }) + expect(results).toStrictEqual([]) + expect(store1.search).not.toHaveBeenCalled() + expect(store2.search).not.toHaveBeenCalled() + }) + }) + + describe('add', () => { + it('writes to all writable stores', async () => { + const store1 = createMockStore('a', { writable: true }) + const store2 = createMockStore('b', { writable: true }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + await mm.add('user likes coffee') + expect(store1.add).toHaveBeenCalledWith('user likes coffee', undefined) + expect(store2.add).toHaveBeenCalledWith('user likes coffee', undefined) + }) + + it('passes metadata to stores', async () => { + const store = createMockStore('a', { writable: true }) + const mm = new MemoryManager({ stores: [store] }) + + await mm.add('fact', { metadata: { source: 'user' } }) + expect(store.add).toHaveBeenCalledWith('fact', { source: 'user' }) + }) + + it('filters to named stores when options.stores is provided', async () => { + const store1 = createMockStore('personal', { writable: true }) + const store2 = createMockStore('team', { writable: true }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + await mm.add('my preference', { stores: ['personal'] }) + expect(store1.add).toHaveBeenCalledWith('my preference', undefined) + expect(store2.add).not.toHaveBeenCalled() + }) + + it('throws when no writable stores match', async () => { + const mm = new MemoryManager({ stores: [createMockStore('a')] }) + await expect(mm.add('fact')).rejects.toThrow('no writable store matched') + }) + + it('throws a not-found error when a named store does not exist', async () => { + const mm = new MemoryManager({ stores: [createMockStore('a', { writable: true })] }) + await expect(mm.add('fact', { stores: ['nonexistent'] })).rejects.toThrow("store 'nonexistent' not found") + }) + + it('throws a read-only error when a named store cannot be written', async () => { + const mm = new MemoryManager({ stores: [createMockStore('readonly')] }) + await expect(mm.add('fact', { stores: ['readonly'] })).rejects.toThrow("store 'readonly' is read-only") + }) + + it('succeeds with partial write failures (some stores fail, some succeed)', async () => { + const store1: MemoryStore = { + name: 'failing', + writable: true, + search: vi.fn().mockResolvedValue([]), + add: vi.fn().mockRejectedValue(new Error('write error')), + } + const store2 = createMockStore('ok', { writable: true }) + const mm = new MemoryManager({ stores: [store1, store2] }) + + await mm.add('fact') + expect(store2.add).toHaveBeenCalledWith('fact', undefined) + }) + + it('throws AggregateError naming the failed stores when all writes fail', async () => { + const store: MemoryStore = { + name: 'failing', + writable: true, + search: vi.fn().mockResolvedValue([]), + add: vi.fn().mockRejectedValue(new Error('write error')), + } + const mm = new MemoryManager({ stores: [store] }) + + await expect(mm.add('fact')).rejects.toThrow('all store writes failed: failing') + }) + }) + + describe('tool store scoping', () => { + function searchTool( + mm: MemoryManager + ): InvokableTool<{ query: string; maxSearchResults?: number; stores?: string[] }, unknown> { + return mm.getTools().find((t) => t.name === 'search_memory') as never + } + + function addTool(mm: MemoryManager): InvokableTool<{ entries: string[]; stores?: string[] }, unknown> { + return mm.getTools().find((t) => t.name === 'add_memory') as never + } + + it('search tool queries all stores when model omits stores', async () => { + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const team = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [personal, team] }) + + await searchTool(mm).invoke({ query: 'q' }) + expect(personal.search).toHaveBeenCalled() + expect(team.search).toHaveBeenCalled() + }) + + it('search tool treats an empty stores array as omitted (searches all)', async () => { + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const team = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [personal, team] }) + + await searchTool(mm).invoke({ query: 'q', stores: [] }) + expect(personal.search).toHaveBeenCalled() + expect(team.search).toHaveBeenCalled() + }) + + it('search tool targets only the requested store when in scope', async () => { + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const team = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [personal, team] }) + + await searchTool(mm).invoke({ query: 'q', stores: ['personal'] }) + expect(personal.search).toHaveBeenCalled() + expect(team.search).not.toHaveBeenCalled() + }) + + it('search tool keeps valid names and warns on out-of-scope names', async () => { + const warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => {}) + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const team = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [personal, team] }) + + await searchTool(mm).invoke({ query: 'q', stores: ['personal', 'nonexistent'] }) + expect(personal.search).toHaveBeenCalled() + expect(team.search).not.toHaveBeenCalled() + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('nonexistent')) + warnSpy.mockRestore() + }) + + it('search tool throws when every requested store is out of scope', async () => { + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const mm = new MemoryManager({ stores: [personal] }) + + await expect(searchTool(mm).invoke({ query: 'q', stores: ['nonexistent'] })).rejects.toThrow( + 'none of the requested memory stores are available' + ) + expect(personal.search).not.toHaveBeenCalled() + }) + + it('add tool writes to all writable stores when model omits stores', async () => { + const personal = createMockStore('personal', { writable: true }) + const team = createMockStore('team', { writable: true }) + const mm = new MemoryManager({ stores: [personal, team], addToolConfig: true }) + + await addTool(mm).invoke({ entries: ['fact'] }) + expect(personal.add).toHaveBeenCalledWith('fact', undefined) + expect(team.add).toHaveBeenCalledWith('fact', undefined) + }) + + it('add tool treats an empty stores array as omitted (writes to all writable)', async () => { + const personal = createMockStore('personal', { writable: true }) + const team = createMockStore('team', { writable: true }) + const mm = new MemoryManager({ stores: [personal, team], addToolConfig: true }) + + await addTool(mm).invoke({ entries: ['fact'], stores: [] }) + expect(personal.add).toHaveBeenCalled() + expect(team.add).toHaveBeenCalled() + }) + + it('add tool excludes read-only stores from its scope', async () => { + const personal = createMockStore('personal', { writable: true }) + const readonly = createMockStore('readonly') + const mm = new MemoryManager({ stores: [personal, readonly], addToolConfig: true }) + + // A read-only store is out of the add tool's scope, so naming it throws. + await expect(addTool(mm).invoke({ entries: ['fact'], stores: ['readonly'] })).rejects.toThrow( + 'none of the requested memory stores are available' + ) + expect(personal.add).not.toHaveBeenCalled() + }) + + it('add tool keeps valid names and warns on out-of-scope names', async () => { + const warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => {}) + const personal = createMockStore('personal', { writable: true }) + const team = createMockStore('team', { writable: true }) + const mm = new MemoryManager({ stores: [personal, team], addToolConfig: true }) + + await addTool(mm).invoke({ entries: ['fact'], stores: ['personal', 'nonexistent'] }) + expect(personal.add).toHaveBeenCalledWith('fact', undefined) + expect(team.add).not.toHaveBeenCalled() + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('nonexistent')) + warnSpy.mockRestore() + }) + + it('add tool throws when every requested store is out of scope', async () => { + const personal = createMockStore('personal', { writable: true }) + const mm = new MemoryManager({ stores: [personal], addToolConfig: true }) + + await expect(addTool(mm).invoke({ entries: ['fact'], stores: ['nonexistent'] })).rejects.toThrow( + 'none of the requested memory stores are available' + ) + expect(personal.add).not.toHaveBeenCalled() + }) + + it('add tool rejects an empty entries array', async () => { + const personal = createMockStore('personal', { writable: true }) + const mm = new MemoryManager({ stores: [personal], addToolConfig: true }) + + await expect(addTool(mm).invoke({ entries: [] })).rejects.toThrow() + expect(personal.add).not.toHaveBeenCalled() + }) + + it('add tool throws when every entry fails to write', async () => { + const failing: MemoryStore = { + name: 'failing', + writable: true, + search: vi.fn().mockResolvedValue([]), + add: vi.fn().mockRejectedValue(new Error('write error')), + } + const mm = new MemoryManager({ stores: [failing], addToolConfig: true }) + + await expect(addTool(mm).invoke({ entries: ['a', 'b'] })).rejects.toThrow('failed to add all 2 entries') + }) + + it('add tool returns counts when some entries succeed and some fail', async () => { + const flaky: MemoryStore = { + name: 'flaky', + writable: true, + search: vi.fn().mockResolvedValue([]), + // First entry's write resolves; second entry's write rejects (its only store), so that entry + // fails entirely while the first succeeds — a genuine per-entry partial outcome. + add: vi.fn().mockResolvedValueOnce(undefined).mockRejectedValueOnce(new Error('write error')), + } + const mm = new MemoryManager({ stores: [flaky], addToolConfig: true }) + + const result = await addTool(mm).invoke({ entries: ['a', 'b'] }) + expect(result).toStrictEqual({ stored: 1, failed: 1 }) + }) + }) + + describe('initAgent', () => { + it('does not throw', () => { + const mm = new MemoryManager({ stores: [createMockStore('test')] }) + expect(() => mm.initAgent({} as any)).not.toThrow() + }) + }) + + describe('injection', () => { + // Builds a fake agent that captures hooks registered against BeforeModelCallEvent, and returns a + // `runInjection` helper that fires the captured callback with the current messages. + function harness(mm: MemoryManager, messages: Message[]) { + const agent = { messages } as unknown as LocalAgent + let injectionCallback: ((event: { agent: LocalAgent }) => unknown) | undefined + agent.addHook = ((eventType: unknown, callback: (event: { agent: LocalAgent }) => unknown) => { + if (eventType === BeforeModelCallEvent) injectionCallback = callback + return () => {} + }) as LocalAgent['addHook'] + + mm.initAgent(agent) + return { + agent, + registered: injectionCallback !== undefined, + runInjection: async () => { + await injectionCallback?.({ agent }) + }, + } + } + + const assistant = (text: string) => new Message({ role: 'assistant', content: [new TextBlock(text)] }) + const user = (text: string) => new Message({ role: 'user', content: [new TextBlock(text)] }) + const injectedTag = (m: Message) => m.metadata?.custom?.['strands:memory-injection'] === true + + it('registers a model-call hook only when injection is enabled', () => { + const off = harness(new MemoryManager({ stores: [createMockStore('s')] }), []) + expect(off.registered).toBe(false) + + const on = harness(new MemoryManager({ stores: [createMockStore('s')], injection: true }), []) + expect(on.registered).toBe(true) + }) + + it('queries with the last assistant text and injects one tagged user message before the last user ask', async () => { + const store = createMockStore('s', { entries: [{ content: 'remembered learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const messages = [user('original task'), assistant('prior step output'), user('next ask')] + const h = harness(mm, messages) + + await h.runInjection() + + expect(store.search).toHaveBeenCalledWith('prior step output', { maxSearchResults: 1 }) + expect(h.agent.messages.map((m) => m.role)).toStrictEqual(['user', 'assistant', 'user', 'user']) + const injected = h.agent.messages[2]! + expect(injectedTag(injected)).toBe(true) + expect((injected.content[0] as TextBlock).text).toContain('remembered learning') + // The user's ask remains the final message. + expect(h.agent.messages[h.agent.messages.length - 1]).toBe(messages[2]) + }) + + it('skips injection when the last message is a tool result', async () => { + const store = createMockStore('s', { entries: [{ content: 'learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const toolResult = new Message({ + role: 'user', + content: [new ToolResultBlock({ toolUseId: 't1', status: 'success', content: [new TextBlock('result')] })], + }) + const messages = [ + user('task'), + new Message({ role: 'assistant', content: [new ToolUseBlock({ name: 'x', toolUseId: 't1', input: {} })] }), + toolResult, + ] + const h = harness(mm, messages) + + await h.runInjection() + + expect(store.search).not.toHaveBeenCalled() + expect(h.agent.messages).toHaveLength(3) + }) + + it('skips the first model call (no prior assistant message)', async () => { + const store = createMockStore('s', { entries: [{ content: 'learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const h = harness(mm, [user('first ask')]) + + await h.runInjection() + + expect(store.search).not.toHaveBeenCalled() + expect(h.agent.messages).toHaveLength(1) + }) + + it('skips when the last assistant message has no text (pure tool use)', async () => { + const store = createMockStore('s', { entries: [{ content: 'learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const messages = [ + new Message({ role: 'assistant', content: [new ToolUseBlock({ name: 'x', toolUseId: 't1', input: {} })] }), + user('ask'), + ] + const h = harness(mm, messages) + + await h.runInjection() + + expect(store.search).not.toHaveBeenCalled() + expect(h.agent.messages).toHaveLength(2) + }) + + it('skips injection when search returns no entries', async () => { + const store = createMockStore('s', { entries: [] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const h = harness(mm, [assistant('prior'), user('ask')]) + + await h.runInjection() + + expect(store.search).toHaveBeenCalled() + expect(h.agent.messages).toHaveLength(2) + }) + + it('strips the prior injection before re-injecting (no accumulation)', async () => { + const store = createMockStore('s', { entries: [{ content: 'learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: true }) + const h = harness(mm, [assistant('prior'), user('ask')]) + + await h.runInjection() + await h.runInjection() + + expect(h.agent.messages.filter(injectedTag)).toHaveLength(1) + expect(h.agent.messages).toHaveLength(3) + }) + + it('honors a custom query and format', async () => { + const store = createMockStore('s', { entries: [{ content: 'A' }, { content: 'B' }] }) + const mm = new MemoryManager({ + stores: [store], + injection: { + maxResults: 2, + query: () => 'custom query', + format: (entries) => entries.map((e) => e.content).join(','), + }, + }) + const h = harness(mm, [assistant('prior'), user('ask')]) + + await h.runInjection() + + expect(store.search).toHaveBeenCalledWith('custom query', { maxSearchResults: 2 }) + expect((h.agent.messages[1]!.content[0] as TextBlock).text).toBe('A,B') + }) + + it('skips when a custom query returns undefined', async () => { + const store = createMockStore('s', { entries: [{ content: 'learning' }] }) + const mm = new MemoryManager({ stores: [store], injection: { query: () => undefined } }) + const h = harness(mm, [assistant('prior'), user('ask')]) + + await h.runInjection() + + expect(store.search).not.toHaveBeenCalled() + expect(h.agent.messages).toHaveLength(2) + }) + }) + + describe('AgentConfig integration', () => { + it('auto-wraps MemoryManagerConfig into MemoryManager instance', () => { + const store = createMockStore('test') + const agent = new Agent({ memoryManager: { stores: [store] } }) + expect(agent.memoryManager).toBeInstanceOf(MemoryManager) + }) + + it('passes through MemoryManager instance unchanged', () => { + const mm = new MemoryManager({ stores: [createMockStore('test')] }) + const agent = new Agent({ memoryManager: mm }) + expect(agent.memoryManager).toBe(mm) + }) + + it('sets memoryManager to undefined when not configured', () => { + const agent = new Agent({}) + expect(agent.memoryManager).toBeUndefined() + }) + }) +}) diff --git a/strands-ts/src/memory/index.ts b/strands-ts/src/memory/index.ts new file mode 100644 index 0000000000..274a945886 --- /dev/null +++ b/strands-ts/src/memory/index.ts @@ -0,0 +1,12 @@ +export { MemoryManager } from './memory-manager.js' +export type { + MemoryEntry, + MemoryStore, + MemoryStoreConfig, + SearchOptions, + SearchMemoryOptions, + AddMemoryOptions, + MemoryToolConfig, + MemoryManagerConfig, + InjectionConfig, +} from './types.js' diff --git a/strands-ts/src/memory/memory-manager.ts b/strands-ts/src/memory/memory-manager.ts new file mode 100644 index 0000000000..8ea55086fb --- /dev/null +++ b/strands-ts/src/memory/memory-manager.ts @@ -0,0 +1,450 @@ +import type { Plugin } from '../plugins/plugin.js' +import type { LocalAgent } from '../types/agent.js' +import type { Tool } from '../tools/tool.js' +import type { + MemoryEntry, + MemoryManagerConfig, + SearchMemoryOptions, + MemoryStore, + AddMemoryOptions, + MemoryToolConfig, + InjectionConfig, +} from './types.js' +import type { JSONValue } from '../types/json.js' +import { tool } from '../tools/tool-factory.js' +import { z } from 'zod' +import { logger } from '../logging/logger.js' +import { BeforeModelCallEvent } from '../hooks/events.js' +import { Message, TextBlock } from '../types/messages.js' + +const SEARCH_TOOL_DESCRIPTION = + 'Search long-term memory for facts, preferences, or context from previous conversations. Use when you need background about the user or topic that may have been discussed before.' + +const ADD_TOOL_DESCRIPTION = + 'Add facts, preferences, or decisions to long-term memory so they are remembered across conversations. Use when the user shares something worth recalling later.' + +const DEFAULT_RESULTS_PER_STORE = 3 + +const DEFAULT_INJECTION_RESULTS = 1 + +/** Marks a message as manager-injected memory context, so it can be stripped before re-injection. */ +const INJECTION_METADATA_KEY = 'strands:memory-injection' + +/** + * Provides cross-session knowledge retrieval and storage for agents. + * + * Manages one or more {@link MemoryStore} backends, exposing `search_memory` and + * `add_memory` tools for agent-driven recall and persistence. Any tools the stores + * themselves provide (via {@link MemoryStore.getTools}) are registered alongside these. + * + * @example + * ```typescript + * import { Agent, MemoryManager } from '@strands-agents/sdk' + * + * // Config shorthand + * const agent = new Agent({ + * model, + * memoryManager: { stores: [myStore], addToolConfig: true }, + * }) + * + * // Class instance (for programmatic access) + * const memoryManager = new MemoryManager({ stores: [myStore], addToolConfig: true }) + * const agent = new Agent({ model, memoryManager }) + * await memoryManager.search('user preferences') + * ``` + */ +export class MemoryManager implements Plugin { + readonly name = 'strands:memory-manager' + private readonly _config: MemoryManagerConfig + private readonly _searchStores: MemoryStore[] + private readonly _addStores: MemoryStore[] + private readonly _searchToolConfig: MemoryToolConfig | false + private readonly _addToolConfig: MemoryToolConfig | false + private readonly _injectionConfig: InjectionConfig | false + + constructor(config: MemoryManagerConfig) { + if (config.stores.length === 0) { + throw new Error('MemoryManager: at least one store is required') + } + + const seenNames = new Set() + for (const store of config.stores) { + if (seenNames.has(store.name)) { + throw new Error(`MemoryManager: duplicate store name '${store.name}'`) + } + seenNames.add(store.name) + + if (store.writable && !store.add) { + throw new Error(`MemoryManager: store '${store.name}' is writable but has no add method`) + } + } + + this._config = config + + // Every store is searchable; the add tool only targets writable stores. + const writableStores = config.stores.filter((s) => s.writable) + + if (config.searchToolConfig === false) { + this._searchToolConfig = false + this._searchStores = [] + } else { + this._searchToolConfig = typeof config.searchToolConfig === 'object' ? config.searchToolConfig : {} + this._searchStores = config.stores + } + + if (config.addToolConfig === undefined || config.addToolConfig === false) { + this._addToolConfig = false + this._addStores = [] + } else { + if (writableStores.length === 0) { + throw new Error('MemoryManager: addToolConfig is enabled but no stores are writable') + } + + this._addToolConfig = typeof config.addToolConfig === 'object' ? config.addToolConfig : {} + this._addStores = writableStores + } + + if (config.injection === undefined || config.injection === false) { + this._injectionConfig = false + } else { + this._injectionConfig = typeof config.injection === 'object' ? config.injection : {} + } + } + + /** + * Initializes the plugin with the agent. + * + * When injection is enabled, registers a {@link BeforeModelCallEvent} hook that searches memory + * and injects results before each model call. Tool registration is handled automatically by the + * PluginRegistry via {@link getTools}. Extraction triggers are deferred to a follow-up PR. + * + * @param agent - The agent this plugin is being attached to + */ + initAgent(agent: LocalAgent): void { + if (this._injectionConfig !== false) { + const injectionConfig = this._injectionConfig + agent.addHook(BeforeModelCallEvent, (event) => this._injectContext(event.agent, injectionConfig)) + } + } + + /** + * Returns tools registered by this plugin. + * + * Includes the manager's own `search_memory` / `add_memory` tools (per their config) plus any + * tools the configured stores expose via {@link MemoryStore.getTools}. + * + * @returns Array of tools to register with the agent + */ + getTools(): Tool[] { + const tools: Tool[] = [] + + if (this._searchToolConfig !== false) { + tools.push(this._createSearchTool(this._searchToolConfig)) + } + + if (this._addToolConfig !== false) { + tools.push(this._createAddTool(this._addToolConfig)) + } + + for (const store of this._config.stores) { + const storeTools = store.getTools?.() ?? [] + tools.push(...storeTools) + } + + return tools + } + + /** + * Search configured stores for entries matching the query. + * + * This method is intentionally unscoped (full access to all configured stores); it is the + * programmatic escape hatch. Tool-level store scoping is applied by the search tool callback. + * When `options.stores` is omitted, all stores are searched; an empty array searches none. + * + * Each store receives `maxSearchResults` individually — results are concatenated in store config + * order. Stores that fail are logged and skipped. + * + * @param query - The search query string + * @param options - Optional max results per-store and store name filter + * @returns Array of memory entries from matching stores + */ + async search(query: string, options?: SearchMemoryOptions): Promise { + logger.debug( + `query=<${query}>, max_search_results=<${options?.maxSearchResults}>, stores=<${options?.stores}> | searching stores` + ) + + const targetStores = + options?.stores !== undefined + ? this._config.stores.filter((s) => options.stores!.includes(s.name)) + : this._config.stores + + if (options?.stores !== undefined && targetStores.length === 0) { + logger.warn(`stores=<${options.stores.join(', ')}> | no stores matched filter`) + } + + const maxSearchResults = options?.maxSearchResults + const settled = await Promise.allSettled( + targetStores.map((store) => + store.search(query, { + maxSearchResults: maxSearchResults ?? store.maxSearchResults ?? DEFAULT_RESULTS_PER_STORE, + }) + ) + ) + + const results: MemoryEntry[] = [] + for (let i = 0; i < settled.length; i++) { + const r = settled[i]! + if (r.status === 'rejected') { + logger.warn(`store=<${targetStores[i]!.name}>, reason=<${r.reason}> | store search failed`) + continue + } + for (const entry of r.value) { + results.push(entry) + } + } + + logger.debug(`results=<${results.length}> | search complete`) + return results + } + + /** + * Add content to writable stores. If `stores` is provided, only writes to those named stores. + * + * This method is intentionally unscoped (full access to all configured writable stores); it is + * the programmatic escape hatch. Tool-level store scoping is applied by the add tool callback. + * When `options.stores` is omitted, all writable stores are targeted; an empty array targets none. + * + * Partial failures are logged. If all writes fail, throws an `AggregateError`. + * + * @param content - The text content to add + * @param options - Optional metadata and store name filter + */ + async add(content: string, options?: AddMemoryOptions): Promise { + let writableStores: MemoryStore[] + + if (options?.stores !== undefined) { + writableStores = options.stores.map((name) => { + const found = this._config.stores.find((s) => s.name === name) + if (!found) { + throw new Error(`MemoryManager: store '${name}' not found`) + } + if (!found.writable) { + throw new Error(`MemoryManager: store '${name}' is read-only`) + } + return found + }) + } else { + writableStores = this._config.stores.filter((s) => s.writable) + } + + if (writableStores.length === 0) { + throw new Error('MemoryManager: no writable store matched') + } + + const settled = await Promise.allSettled(writableStores.map((s) => s.add!(content, options?.metadata))) + + const failures: { store: string; reason: unknown }[] = [] + for (let i = 0; i < settled.length; i++) { + const r = settled[i]! + if (r.status === 'rejected') { + const storeName = writableStores[i]!.name + logger.warn(`store=<${storeName}>, reason=<${r.reason}> | store write failed`) + failures.push({ store: storeName, reason: r.reason }) + } + } + if (failures.length === writableStores.length) { + throw new AggregateError( + failures.map((f) => f.reason), + `MemoryManager: all store writes failed: ${failures.map((f) => f.store).join(', ')}` + ) + } + } + + /** + * Searches memory and injects the results as a `user` message before the latest user message. + * + * Runs before each model call. Previously injected messages are stripped first so only the + * current results are present. Injection is skipped unless the latest message is a user ask + * (not a tool result), which keeps the user's ask the final message the model sees. + * + * @param agent - The agent whose messages are being injected into + * @param config - The resolved injection configuration + */ + private async _injectContext(agent: LocalAgent, config: InjectionConfig): Promise { + // Strip prior injections so only the latest results are present. + agent.messages = agent.messages.filter((m) => m.metadata?.custom?.[INJECTION_METADATA_KEY] !== true) + + // Only inject when the latest message is a user ask, not a tool result, so the ask stays last. + const last = agent.messages[agent.messages.length - 1] + if (!last || last.role !== 'user' || last.content.some((b) => b.type === 'toolResultBlock')) { + return + } + + const query = this._resolveInjectionQuery(agent.messages, config) + if (!query?.trim()) { + return + } + + const maxSearchResults = config.maxResults ?? DEFAULT_INJECTION_RESULTS + const entries = (await this.search(query, { maxSearchResults })).slice(0, maxSearchResults) + if (entries.length === 0) { + return + } + + const text = (config.format ?? this._defaultInjectionFormat)(entries) + const message = new Message({ + role: 'user', + content: [new TextBlock(text)], + metadata: { custom: { [INJECTION_METADATA_KEY]: true } }, + }) + + // Insert immediately before the latest user message. + agent.messages.splice(agent.messages.length - 1, 0, message) + } + + /** + * Derives the injection search query. Uses the configured `query` if provided, otherwise the text + * of the most recent assistant message. + * + * @returns The query string, or undefined when none is available + */ + private _resolveInjectionQuery(messages: Message[], config: InjectionConfig): string | undefined { + if (config.query) { + return config.query(messages.map((m) => m.toJSON())) + } + + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]! + if (message.role !== 'assistant') { + continue + } + const text = message.content + .filter((b): b is TextBlock => b.type === 'textBlock') + .map((b) => b.text) + .join('\n') + .trim() + return text.length > 0 ? text : undefined + } + + return undefined + } + + /** Default injection format: wraps entry contents in an XML block. */ + private _defaultInjectionFormat(entries: MemoryEntry[]): string { + return `\n${entries.map((e) => e.content).join('\n')}\n` + } + + /** + * Resolves the store names a tool callback should target against the tool's scoped set. + * + * - Omitting `requested` (or passing an empty array) targets all scoped stores. + * - Names that are in scope are kept; out-of-scope names are dropped with a warning. + * - When every requested name is out of scope, throws so the model receives an actionable error + * (the tool layer turns the thrown error into a model-visible result it can correct from). + * + * @param scopedNames - Store names available to this tool + * @param requested - Store names the model asked for, if any + * @returns A non-empty list of in-scope store names to target + */ + private _resolveToolTargets(scopedNames: string[], requested?: string[]): string[] { + if (requested === undefined || requested.length === 0) { + return scopedNames + } + + const inScope = requested.filter((name) => scopedNames.includes(name)) + const outOfScope = requested.filter((name) => !scopedNames.includes(name)) + + if (inScope.length === 0) { + throw new Error( + `MemoryManager: requested=<${requested.join(', ')}> | none of the requested memory stores are available; available stores: ${scopedNames.join(', ')}` + ) + } + + if (outOfScope.length > 0) { + logger.warn(`requested=<${outOfScope.join(', ')}> | ignoring memory stores outside this tool's scope`) + } + + return inScope + } + + private _createSearchTool(config: MemoryToolConfig): Tool { + let description = config.description ?? SEARCH_TOOL_DESCRIPTION + const storeDescriptions = this._searchStores + .filter((s) => s.description) + .map((s) => `- ${s.name}: ${s.description}`) + if (storeDescriptions.length > 0) { + description += `\n\nAvailable memory stores:\n${storeDescriptions.join('\n')}` + description += + '\n\nYou can target one or more memory stores by name if you know which domains are relevant, or omit the stores parameter to search all.' + } + + const scopedNames = this._searchStores.map((s) => s.name) + + const inputSchema = z.object({ + query: z.string().describe('What to search for'), + maxSearchResults: z.number().optional().describe('Maximum number of results per store'), + stores: z + .array(z.string()) + .optional() + .describe('Filter to specific stores by name. Omit to search all available stores.'), + }) + + return tool({ + name: config.name ?? 'search_memory', + description, + inputSchema, + callback: async (input) => { + const stores = this._resolveToolTargets(scopedNames, input.stores) + const results = await this.search(input.query, { + ...(input.maxSearchResults != null && { maxSearchResults: input.maxSearchResults }), + stores, + }) + return results.map((entry) => ({ + content: entry.content, + ...(entry.metadata && { metadata: entry.metadata }), + })) as JSONValue + }, + }) + } + + private _createAddTool(config: MemoryToolConfig): Tool { + let description = config.description ?? ADD_TOOL_DESCRIPTION + const storeDescriptions = this._addStores.filter((s) => s.description).map((s) => `- ${s.name}: ${s.description}`) + if (storeDescriptions.length > 0) { + description += `\n\nAvailable writable stores:\n${storeDescriptions.join('\n')}` + description += + '\n\nYou can target a specific store by name to route facts to the right place, or omit to add to all writable stores.' + } + + const scopedNames = this._addStores.map((s) => s.name) + + const inputSchema = z.object({ + entries: z.array(z.string()).min(1).describe('Data to add to long-term memory'), + stores: z + .array(z.string()) + .optional() + .describe('Target specific stores by name. Omit to add to all writable stores.'), + }) + + return tool({ + name: config.name ?? 'add_memory', + description, + inputSchema, + callback: async (input) => { + const stores = this._resolveToolTargets(scopedNames, input.stores) + const settled = await Promise.allSettled(input.entries.map((content) => this.add(content, { stores }))) + const stored = settled.filter((r) => r.status === 'fulfilled').length + const failures = settled.filter((r) => r.status === 'rejected') as PromiseRejectedResult[] + + if (stored === 0 && failures.length > 0) { + throw new AggregateError( + failures.map((f) => f.reason), + `MemoryManager: failed to add all ${failures.length} entries` + ) + } + + return { stored, failed: failures.length } as JSONValue + }, + }) + } +} diff --git a/strands-ts/src/memory/stores/bedrock-knowledge-base-store.ts b/strands-ts/src/memory/stores/bedrock-knowledge-base-store.ts new file mode 100644 index 0000000000..182c20c1c0 --- /dev/null +++ b/strands-ts/src/memory/stores/bedrock-knowledge-base-store.ts @@ -0,0 +1,194 @@ +import { + BedrockAgentRuntimeClient, + type BedrockAgentRuntimeClientConfig, + RetrieveCommand, + type RetrievalFilter, +} from '@aws-sdk/client-bedrock-agent-runtime' +import { + BedrockAgentClient, + type BedrockAgentClientConfig, + IngestKnowledgeBaseDocumentsCommand, +} from '@aws-sdk/client-bedrock-agent' +import { v7 as uuidv7 } from 'uuid' + +import type { MemoryEntry, MemoryStore, MemoryStoreConfig, SearchOptions } from '../types.js' +import type { JSONValue } from '../../types/json.js' + +export interface BedrockKnowledgeBaseStoreConfig extends MemoryStoreConfig { + knowledgeBaseId: string + /** + * Data source to ingest into when writing. Required for `add` to succeed — without it, write + * calls throw, since the knowledge base has no destination to ingest into. + */ + dataSourceId?: string + scope?: string + scopeMetadataKey?: string + filter?: RetrievalFilter + runtimeClientConfig?: BedrockAgentRuntimeClientConfig + runtimeClient?: BedrockAgentRuntimeClient + agentClientConfig?: BedrockAgentClientConfig + agentClient?: BedrockAgentClient +} + +export class BedrockKnowledgeBaseStore implements MemoryStore { + readonly name: string + readonly description?: string + readonly maxSearchResults?: number + readonly writable: boolean + + private readonly _runtimeClient: BedrockAgentRuntimeClient + private _agentClient: BedrockAgentClient | undefined + private readonly _agentClientConfig: BedrockAgentClientConfig | undefined + private readonly _knowledgeBaseId: string + private readonly _dataSourceId: string | undefined + private readonly _scope: string | undefined + private readonly _scopeMetadataKey: string + private readonly _filter: RetrievalFilter | undefined + + constructor(config: BedrockKnowledgeBaseStoreConfig) { + this.name = config.name + if (config.description !== undefined) this.description = config.description + if (config.maxSearchResults !== undefined) this.maxSearchResults = config.maxSearchResults + this.writable = config.writable ?? false + + this._runtimeClient = config.runtimeClient ?? new BedrockAgentRuntimeClient(config.runtimeClientConfig ?? {}) + this._agentClient = config.agentClient + this._agentClientConfig = config.agentClientConfig + this._knowledgeBaseId = config.knowledgeBaseId + this._dataSourceId = config.dataSourceId + this._scope = config.scope + this._scopeMetadataKey = config.scopeMetadataKey ?? 'namespace' + + if (config.filter) { + this._filter = config.filter + } else if (config.scope) { + this._filter = { + equals: { + key: this._scopeMetadataKey, + value: config.scope, + }, + } + } + } + + async search(query: string, options?: SearchOptions): Promise { + const limit = options?.maxSearchResults ?? this.maxSearchResults ?? 10 + + const response = await this._runtimeClient.send( + new RetrieveCommand({ + knowledgeBaseId: this._knowledgeBaseId, + retrievalQuery: { text: query }, + retrievalConfiguration: { + vectorSearchConfiguration: { + numberOfResults: limit, + ...(this._filter && { filter: this._filter }), + }, + }, + }) + ) + + return (response.retrievalResults ?? []).map((result) => { + const metadata: Record = {} + if (result.metadata) { + for (const [key, value] of Object.entries(result.metadata)) { + metadata[key] = value as JSONValue + } + } + if (result.location) { + metadata._location = result.location as unknown as JSONValue + } + if (result.score != null) { + metadata.score = result.score + } + + return { + content: result.content?.text ?? '', + metadata, + } + }) + } + + async add(content: string, metadata?: Record): Promise { + const dataSourceId = this._requireDataSourceId() + const id = uuidv7() + + const inlineAttributes: Array<{ + key: string + value: + | { type: 'STRING'; stringValue: string } + | { type: 'NUMBER'; numberValue: number } + | { type: 'BOOLEAN'; booleanValue: boolean } + }> = [] + + if (this._scope) { + inlineAttributes.push({ + key: this._scopeMetadataKey, + value: { type: 'STRING' as const, stringValue: this._scope }, + }) + } + + if (metadata) { + for (const [key, value] of Object.entries(metadata)) { + if (typeof value === 'string') { + inlineAttributes.push({ + key, + value: { type: 'STRING' as const, stringValue: value }, + }) + } else if (typeof value === 'number') { + inlineAttributes.push({ + key, + value: { type: 'NUMBER' as const, numberValue: value }, + }) + } else if (typeof value === 'boolean') { + inlineAttributes.push({ + key, + value: { type: 'BOOLEAN' as const, booleanValue: value }, + }) + } + } + } + + await this._getAgentClient().send( + new IngestKnowledgeBaseDocumentsCommand({ + knowledgeBaseId: this._knowledgeBaseId, + dataSourceId, + documents: [ + { + content: { + dataSourceType: 'CUSTOM', + custom: { + customDocumentIdentifier: { id }, + sourceType: 'IN_LINE', + inlineContent: { + type: 'TEXT', + textContent: { data: content }, + }, + }, + }, + metadata: { + type: 'IN_LINE_ATTRIBUTE', + inlineAttributes, + }, + }, + ], + }) + ) + } + + private _requireDataSourceId(): string { + if (!this._dataSourceId) { + throw new Error( + 'BedrockKnowledgeBaseStore: dataSourceId is required for write operations. ' + + 'Provide it in the config to enable add().' + ) + } + return this._dataSourceId + } + + private _getAgentClient(): BedrockAgentClient { + if (!this._agentClient) { + this._agentClient = new BedrockAgentClient(this._agentClientConfig ?? {}) + } + return this._agentClient + } +} diff --git a/strands-ts/src/memory/types.ts b/strands-ts/src/memory/types.ts new file mode 100644 index 0000000000..848c0fa8c6 --- /dev/null +++ b/strands-ts/src/memory/types.ts @@ -0,0 +1,152 @@ +import type { JSONValue } from '../types/json.js' +import type { Tool } from '../tools/tool.js' +import type { MessageData } from '../types/messages.js' + +/** + * A single entry retrieved from or stored to a memory store. + */ +export interface MemoryEntry { + /** The textual content of this memory entry. */ + content: string + /** Optional metadata (e.g., score, source, id, timestamp). */ + metadata?: Record +} + +/** + * Options passed to {@link MemoryStore.search}. + * Store implementations may extend this with additional fields in their own signatures. + */ +export interface SearchOptions { + /** Maximum number of results to return. */ + maxSearchResults?: number +} + +/** + * Common configuration shared by all built-in memory stores. + * + * Store implementations should extend this so identity, result limits, and writability are + * consistent across every store. Concrete stores add their own backend-specific fields. + */ +export interface MemoryStoreConfig { + /** Identifier for this store, used to target specific stores in search/add tools. */ + name: string + /** Human-readable description of what this store contains. Included in tool descriptions. */ + description?: string + /** + * Default maximum number of results this store returns per search, used when a caller does not + * pass a per-call `maxSearchResults`. Defaults to 3. + */ + maxSearchResults?: number + /** + * Whether this store instance accepts writes. Concrete stores resolve this to a definite + * `writable` boolean on the {@link MemoryStore} (defaulting to `false` when omitted). + * + * @defaultValue false + */ + writable?: boolean +} + +/** + * Interface for a memory store backend. + * + * Every store is searchable. The `writable` flag declares whether the store also accepts writes, + * which is how the {@link MemoryManager} decides where to route them: `search_memory` queries all + * stores, while `add_memory` only writes to `writable` stores. + */ +export interface MemoryStore { + /** Identifier for this store, used to target specific stores in search/add tools. */ + readonly name: string + /** + * Whether this store accepts writes. + * - `false`: searchable only; never written to. + * - `true`: searchable and writable. Requires `add` to be implemented. + */ + readonly writable: boolean + /** Human-readable description of what this store contains. Included in tool descriptions. */ + readonly description?: string + /** Default max results per query for this store. Defaults to 3. */ + readonly maxSearchResults?: number + /** Search the store for entries matching the query, ordered by relevance. */ + search(query: string, options?: SearchOptions): Promise + /** + * Add content to the store. Required when `writable` is `true`; ignored otherwise. + * A store may implement `add` while declaring `writable: false`, in which case it is never invoked. + */ + add?(content: string, metadata?: Record): Promise + /** + * Returns store-specific tools to register with the agent, alongside the manager's own + * `search_memory` / `add_memory` tools. Implement to expose backend-specific capabilities + * (e.g. a store-native query tool). Optional — mirrors {@link Plugin.getTools}. + * + * @returns Array of tools provided by this store + */ + getTools?(): Tool[] +} + +/** + * Options for {@link MemoryManager.search}. + */ +export interface SearchMemoryOptions { + /** Maximum number of results per store. */ + maxSearchResults?: number + /** Filter to specific stores by name. Omit to search all. */ + stores?: string[] +} + +/** + * Options for {@link MemoryManager.add}. + */ +export interface AddMemoryOptions { + /** Metadata to associate with the added entry. */ + metadata?: Record + /** Filter to specific writable stores by name. Omit to write to all. */ + stores?: string[] +} + +/** + * Configuration for customizing a memory tool's name or description. + * + * Store targeting is derived from each store's `writable` flag (see {@link MemoryStore}), not + * configured here: `search_memory` targets all stores, `add_memory` targets `writable` stores. + */ +export interface MemoryToolConfig { + /** Custom tool name. */ + name?: string + /** Custom tool description. */ + description?: string +} + +/** + * Configuration for passive context injection. + * + * When enabled, the {@link MemoryManager} searches memory before each model call and injects the + * top results as a `user` message placed immediately before the user's latest message, so relevant + * knowledge is always present without the model choosing to search. Injection only runs when the + * latest message is a user ask (not a tool result), keeping the user's ask the final message the + * model sees. + */ +export interface InjectionConfig { + /** Maximum number of entries to retrieve and inject per model call. Defaults to 1. */ + maxResults?: number + /** + * Derives the search query from the current conversation. Return `undefined` or an empty string + * to skip injection for this call. Defaults to the text of the most recent assistant message. + */ + query?: (messages: MessageData[]) => string | undefined + /** Renders retrieved entries into the injected message text. Defaults to an XML block. */ + format?: (entries: MemoryEntry[]) => string +} + +/** + * Configuration for the {@link MemoryManager}. + */ +export interface MemoryManagerConfig { + /** One or more memory stores to manage. */ + stores: MemoryStore[] + /** Search tool configuration. Defaults to `true` (auto-created targeting all stores). */ + searchToolConfig?: MemoryToolConfig | boolean + /** Add tool configuration. Defaults to `false` (opt-in). */ + addToolConfig?: MemoryToolConfig | boolean + /** Passive context injection. Defaults to `false` (opt-in). `true` uses default injection settings. */ + injection?: boolean | InjectionConfig +}