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..16a7679cbb 100644 --- a/strands-ts/src/index.ts +++ b/strands-ts/src/index.ts @@ -286,6 +286,18 @@ 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, + SearchOptions, + MemorySearchOptions, + MemoryStoreOptions, + MemoryToolConfig, + MemoryManagerConfig, +} 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..40849b7fad --- /dev/null +++ b/strands-ts/src/memory/__tests__/memory-manager.test.ts @@ -0,0 +1,360 @@ +import { describe, it, expect, vi } from 'vitest' +import { Agent } from '../../agent/agent.js' +import { MemoryManager } from '../memory-manager.js' +import type { MemoryStore, MemoryEntry } from '../types.js' +import type { InvokableTool } from '../../tools/tool.js' + +function createMockStore( + name: string, + options?: { entries?: MemoryEntry[]; writable?: boolean; description?: string; limit?: number } +): MemoryStore { + const store: MemoryStore = { + name, + ...(options?.description && { description: options.description }), + ...(options?.limit != null && { limit: options.limit }), + search: vi.fn().mockResolvedValue(options?.entries ?? []), + } + if (options?.writable) { + store.add = vi.fn().mockResolvedValue(undefined) + } + return store +} + +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 storeToolConfig references non-existent store', () => { + expect( + () => + new MemoryManager({ + stores: [createMockStore('a')], + storeToolConfig: { stores: ['nonexistent'] }, + }) + ).toThrow("store 'nonexistent' not found") + }) + + it('throws when storeToolConfig targets no writable stores', () => { + expect( + () => + new MemoryManager({ + stores: [createMockStore('a')], + storeToolConfig: true, + }) + ).toThrow('storeToolConfig targets no writable stores') + }) + + it('throws when storeToolConfig is true with multiple writable stores and no explicit stores', () => { + expect( + () => + new MemoryManager({ + stores: [createMockStore('a', { writable: true }), createMockStore('b', { writable: true })], + storeToolConfig: true, + }) + ).toThrow('must specify `stores` when multiple writable stores are configured') + }) + + it('allows storeToolConfig true with single writable store', () => { + const mm = new MemoryManager({ + stores: [createMockStore('a', { writable: true })], + storeToolConfig: true, + }) + expect(mm.getTools().map((t) => t.name)).toContain('store_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 store tool when storeToolConfig is enabled', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + storeToolConfig: true, + }) + const tools = mm.getTools() + expect(tools.map((t) => t.name)).toStrictEqual(['search_memory', 'store_memory']) + }) + + it('does not register store 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 storeToolConfig is false', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + searchToolConfig: false, + storeToolConfig: false, + }) + expect(mm.getTools()).toStrictEqual([]) + }) + + it('uses custom tool names from MemoryToolConfig', () => { + const mm = new MemoryManager({ + stores: [createMockStore('test', { writable: true })], + searchToolConfig: { name: 'recall' }, + storeToolConfig: { name: 'remember', stores: ['test'] }, + }) + 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 store tool description', () => { + const store = createMockStore('notes', { writable: true, description: 'Personal notes' }) + const mm = new MemoryManager({ stores: [store], storeToolConfig: true }) + const tools = mm.getTools() + const storeTool = tools.find((t) => t.name === 'store_memory')! + expect(storeTool.description).toContain('notes: Personal notes') + expect(storeTool.description).toContain('target a specific store by name') + }) + }) + + 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 limit to each store', async () => { + const store = createMockStore('a', { limit: 5 }) + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query') + expect(store.search).toHaveBeenCalledWith('query', { limit: 5 }) + }) + + it('overrides per-store limit with options.limit', async () => { + const store = createMockStore('a', { limit: 5 }) + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query', { limit: 2 }) + expect(store.search).toHaveBeenCalledWith('query', { limit: 2 }) + }) + + it('defaults to limit of 3 when no limit configured', async () => { + const store = createMockStore('a') + const mm = new MemoryManager({ stores: [store] }) + + await mm.search('query') + expect(store.search).toHaveBeenCalledWith('query', { limit: 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', 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('store', () => { + 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.store('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.store('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.store('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.store('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.store('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.store('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', + 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.store('fact') + expect(store2.add).toHaveBeenCalledWith('fact', undefined) + }) + + it('throws AggregateError naming the failed stores when all writes fail', async () => { + const store: MemoryStore = { + name: 'failing', + search: vi.fn().mockResolvedValue([]), + add: vi.fn().mockRejectedValue(new Error('write error')), + } + const mm = new MemoryManager({ stores: [store] }) + + await expect(mm.store('fact')).rejects.toThrow('all store writes failed: failing') + }) + }) + + describe('tool store scoping', () => { + function searchTool( + mm: MemoryManager + ): InvokableTool<{ query: string; limit?: number; stores?: string[] }, unknown> { + return mm.getTools().find((t) => t.name === 'search_memory') as never + } + + function storeTool(mm: MemoryManager): InvokableTool<{ entries: string[]; stores?: string[] }, unknown> { + return mm.getTools().find((t) => t.name === 'store_memory') as never + } + + it('search tool only queries scoped 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], searchToolConfig: { stores: ['personal'] } }) + + await searchTool(mm).invoke({ query: 'q' }) + expect(personal.search).toHaveBeenCalled() + expect(team.search).not.toHaveBeenCalled() + }) + + it('search tool drops out-of-scope store names supplied by the model', async () => { + const personal = createMockStore('personal', { entries: [{ content: 'personal fact' }] }) + const team = createMockStore('team', { entries: [{ content: 'team fact' }] }) + const mm = new MemoryManager({ stores: [personal, team], searchToolConfig: { stores: ['personal'] } }) + + await searchTool(mm).invoke({ query: 'q', stores: ['team'] }) + expect(personal.search).not.toHaveBeenCalled() + expect(team.search).not.toHaveBeenCalled() + }) + + it('store tool only writes to scoped 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], storeToolConfig: { stores: ['personal'] } }) + + await storeTool(mm).invoke({ entries: ['fact'] }) + expect(personal.add).toHaveBeenCalledWith('fact', undefined) + expect(team.add).not.toHaveBeenCalled() + }) + + it('store tool drops out-of-scope store names supplied by the model', async () => { + const personal = createMockStore('personal', { writable: true }) + const team = createMockStore('team', { writable: true }) + const mm = new MemoryManager({ stores: [personal, team], storeToolConfig: { stores: ['personal'] } }) + + const result = await storeTool(mm).invoke({ entries: ['fact'], stores: ['team'] }) + expect(personal.add).not.toHaveBeenCalled() + expect(team.add).not.toHaveBeenCalled() + expect(result).toStrictEqual({ stored: 0, failed: 1 }) + }) + }) + + describe('initAgent', () => { + it('does not throw', () => { + const mm = new MemoryManager({ stores: [createMockStore('test')] }) + expect(() => mm.initAgent({} as any)).not.toThrow() + }) + }) + + 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..d9b54b8926 --- /dev/null +++ b/strands-ts/src/memory/index.ts @@ -0,0 +1,11 @@ +export { MemoryManager } from './memory-manager.js' +export type { + MemoryEntry, + MemoryStore, + MemoryStoreConfig, + SearchOptions, + MemorySearchOptions, + MemoryStoreOptions, + MemoryToolConfig, + MemoryManagerConfig, +} 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..e01a5f3f32 --- /dev/null +++ b/strands-ts/src/memory/memory-manager.ts @@ -0,0 +1,321 @@ +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, + MemorySearchOptions, + MemoryStore, + MemoryStoreOptions, + MemoryToolConfig, +} 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' + +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 STORE_TOOL_DESCRIPTION = + 'Store facts, preferences, or decisions that should be remembered across conversations. Use when the user shares something worth recalling later.' + +const DEFAULT_RESULTS_PER_STORE = 3 + +/** + * Provides cross-session knowledge retrieval and storage for agents. + * + * Manages one or more {@link MemoryStore} backends, exposing `search_memory` and + * `store_memory` tools for agent-driven recall and persistence. + * + * @example + * ```typescript + * import { Agent, MemoryManager } from '@strands-agents/sdk' + * + * // Config shorthand + * const agent = new Agent({ + * model, + * memoryManager: { stores: [myStore], storeToolConfig: true }, + * }) + * + * // Class instance (for programmatic access) + * const memoryManager = new MemoryManager({ stores: [myStore], storeToolConfig: 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 _storeStores: MemoryStore[] + private readonly _searchToolConfig: MemoryToolConfig | false + private readonly _storeToolConfig: MemoryToolConfig | 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) + } + + this._config = config + + if (config.searchToolConfig === false) { + this._searchToolConfig = false + this._searchStores = [] + } else { + const toolConfig = typeof config.searchToolConfig === 'object' ? config.searchToolConfig : {} + this._searchStores = this._resolveStores(config.stores, toolConfig.stores) + this._searchToolConfig = toolConfig + } + + if (config.storeToolConfig === undefined || config.storeToolConfig === false) { + this._storeToolConfig = false + this._storeStores = [] + } else { + const toolConfig = typeof config.storeToolConfig === 'object' ? config.storeToolConfig : {} + const resolved = this._resolveStores(config.stores, toolConfig.stores).filter((s) => s.add) + + if (resolved.length === 0) { + throw new Error('MemoryManager: storeToolConfig targets no writable stores') + } + + if (config.storeToolConfig === true && resolved.length > 1 && !toolConfig.stores) { + throw new Error( + 'MemoryManager: storeToolConfig must specify `stores` when multiple writable stores are configured' + ) + } + + this._storeStores = resolved + this._storeToolConfig = toolConfig + } + } + + /** + * Initializes the plugin with the agent. + * + * No lifecycle hooks are registered in this version; context injection and extraction + * triggers are deferred to a follow-up PR. Tool registration is handled automatically + * by the PluginRegistry via {@link getTools}. + * + * @param _agent - The agent this plugin is being attached to + */ + initAgent(_agent: LocalAgent): void {} + + /** + * Returns tools registered by this plugin. + * + * @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._storeToolConfig !== false) { + tools.push(this._createStoreTool(this._storeToolConfig)) + } + + 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 the `limit` individually — results are concatenated in store config order. + * Stores that fail are logged and skipped. + * + * @param query - The search query string + * @param options - Optional limit per-store and store name filter + * @returns Array of memory entries from matching stores + */ + async search(query: string, options?: MemorySearchOptions): Promise { + logger.debug(`query=<${query}>, limit=<${options?.limit}>, 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 limit = options?.limit + const settled = await Promise.allSettled( + targetStores.map((store) => + store.search(query, { limit: limit ?? 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 + } + + /** + * Store content in 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 store 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 store + * @param options - Optional metadata and store name filter + */ + async store(content: string, options?: MemoryStoreOptions): 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.add) { + throw new Error(`MemoryManager: store '${name}' is read-only`) + } + return found + }) + } else { + writableStores = this._config.stores.filter((s) => s.add) + } + + 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(', ')}` + ) + } + } + + private _resolveStores(allStores: MemoryStore[], scoped?: string[]): MemoryStore[] { + if (!scoped || scoped.length === 0) return allStores + + return scoped.map((name) => { + const found = allStores.find((s) => s.name === name) + if (!found) { + throw new Error(`MemoryManager: store '${name}' not found`) + } + return found + }) + } + + 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'), + limit: z.number().optional().describe('Maximum number of results per store'), + stores: z + .array(z.string()) + .optional() + .describe( + 'Filter to specific stores by name. Omit entirely to search all available stores (do not pass an empty array).' + ), + }) + + return tool({ + name: config.name ?? 'search_memory', + description, + inputSchema, + callback: async (input) => { + const stores = input.stores != null ? input.stores.filter((name) => scopedNames.includes(name)) : scopedNames + const results = await this.search(input.query, { + ...(input.limit != null && { limit: input.limit }), + stores, + }) + return results.map((entry) => ({ + content: entry.content, + ...(entry.metadata && { metadata: entry.metadata }), + })) as JSONValue + }, + }) + } + + private _createStoreTool(config: MemoryToolConfig): Tool { + let description = config.description ?? STORE_TOOL_DESCRIPTION + const storeDescriptions = this._storeStores.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 store in all writable stores.' + } + + const scopedNames = this._storeStores.map((s) => s.name) + + const inputSchema = z.object({ + entries: z.array(z.string()).describe('Data to store in long-term memory'), + stores: z + .array(z.string()) + .optional() + .describe( + 'Target specific stores by name. Omit entirely to store in all writable stores (do not pass an empty array).' + ), + }) + + return tool({ + name: config.name ?? 'store_memory', + description, + inputSchema, + callback: async (input) => { + const stores = input.stores != null ? input.stores.filter((name) => scopedNames.includes(name)) : scopedNames + const settled = await Promise.allSettled(input.entries.map((content) => this.store(content, { stores }))) + const stored = settled.filter((r) => r.status === 'fulfilled').length + const failed = settled.filter((r) => r.status === 'rejected').length + return { stored, failed } 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..921d71c50a --- /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?.limit ?? 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..4379396d8c --- /dev/null +++ b/strands-ts/src/memory/types.ts @@ -0,0 +1,123 @@ +import type { JSONValue } from '../types/json.js' +import type { Tool } from '../tools/tool.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. */ + limit?: 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/store 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 `limit`. Defaults to 3. + */ + maxSearchResults?: number + /** + * Whether the caller wants this store instance to accept writes. This is per-instance intent, + * not a capability of the store type: a backend may support writing, yet you can pin a given + * instance to read-only by leaving this unset. + * + * @defaultValue false + */ + writable?: boolean +} + +/** + * Interface for a memory store backend. + * + * Only `search` is required. Stores the caller wants to write to additionally implement `add` and + * report `writable: true`. + */ +export interface MemoryStore { + /** Identifier for this store, used to target specific stores in search/store tools. */ + readonly name: string + /** Human-readable description of what this store contains. Included in tool descriptions. */ + readonly description?: string + /** Default maximum number of results this store returns per search. Defaults to 3. */ + readonly maxSearchResults?: number + /** + * Whether this instance accepts writes, reflecting the caller's per-instance intent. A store the + * caller made writable exposes `add` and reports `writable: true`; a read-only instance omits + * `add`. When omitted, writability is inferred from the presence of `add` (backwards-compatible + * default for ad-hoc stores). + */ + readonly writable?: boolean + /** Search the store for entries matching the query, ordered by relevance. */ + search(query: string, options?: SearchOptions): Promise + /** Add content to the store. Optional — only present on stores the caller made writable. */ + add?(content: string, metadata?: Record): Promise + /** + * Returns store-specific tools to register with the agent. Optional — implement to expose + * backend-specific capabilities (e.g. management or query tools) beyond the manager's + * `search`/`store` tools. + */ + getTools?(): Tool[] +} + +/** + * Options for {@link MemoryManager.search}. + */ +export interface MemorySearchOptions { + /** Maximum number of results per store. */ + limit?: number + /** Filter to specific stores by name. Omit to search all. */ + stores?: string[] +} + +/** + * Options for {@link MemoryManager.store}. + */ +export interface MemoryStoreOptions { + /** Metadata to associate with the stored 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, description, or store scoping. + */ +export interface MemoryToolConfig { + /** Custom tool name. */ + name?: string + /** Custom tool description. */ + description?: string + /** Scopes which stores this tool targets by name. Defaults to all applicable stores. */ + stores?: 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 + /** Store tool configuration. Defaults to `false` (opt-in). */ + storeToolConfig?: MemoryToolConfig | boolean +}