diff --git a/typescript/tests/agents/BedrockFlowsAgent.test.ts b/typescript/tests/agents/BedrockFlowsAgent.test.ts new file mode 100644 index 000000000..be777128e --- /dev/null +++ b/typescript/tests/agents/BedrockFlowsAgent.test.ts @@ -0,0 +1,119 @@ +import { BedrockFlowsAgent, BedrockFlowsAgentOptions } from '../../src/agents/bedrockFlowsAgent'; +import { BedrockAgentRuntimeClient, InvokeFlowCommand } from '@aws-sdk/client-bedrock-agent-runtime'; +import { ConversationMessage, ParticipantRole } from '../../src/types'; + +jest.mock('@aws-sdk/client-bedrock-agent-runtime'); +jest.mock('../../src/common/src/awsSdkUtils', () => ({ + addUserAgentMiddleware: jest.fn(), +})); + +describe('BedrockFlowsAgent', () => { + let mockClient: jest.Mocked; + + const defaultOptions: BedrockFlowsAgentOptions = { + name: 'TestFlowsAgent', + description: 'Test BedrockFlowsAgent', + flowIdentifier: 'flow-123', + flowAliasIdentifier: 'alias-456', + }; + + function makeResponseStream(document: string) { + async function* gen() { + yield { flowOutputEvent: { content: { document } } }; + } + return { responseStream: gen() }; + } + + beforeEach(() => { + jest.clearAllMocks(); + mockClient = { send: jest.fn(), use: jest.fn(), destroy: jest.fn() } as any; + (BedrockAgentRuntimeClient as jest.Mock).mockImplementation(() => mockClient); + (InvokeFlowCommand as unknown as jest.Mock).mockImplementation((p) => p); + }); + + describe('constructor', () => { + it('creates a BedrockAgentRuntimeClient with no arguments when no region given', () => { + new BedrockFlowsAgent(defaultOptions); + expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith(); + }); + + it('creates a BedrockAgentRuntimeClient with specified region', () => { + new BedrockFlowsAgent({ ...defaultOptions, region: 'eu-west-1' }); + expect(BedrockAgentRuntimeClient).toHaveBeenCalledWith({ region: 'eu-west-1' }); + }); + + it('uses a pre-built client when bedrockAgentClient is provided', () => { + const prebuilt = { send: jest.fn() } as any; + const agent = new BedrockFlowsAgent({ ...defaultOptions, bedrockAgentClient: prebuilt }); + expect(agent).toBeDefined(); + expect(BedrockAgentRuntimeClient).not.toHaveBeenCalled(); + }); + + it('defaults enableTrace to false', () => { + const agent = new BedrockFlowsAgent(defaultOptions); + expect((agent as any).enableTrace).toBe(false); + }); + + it('sets enableTrace when provided', () => { + const agent = new BedrockFlowsAgent({ ...defaultOptions, enableTrace: true }); + expect((agent as any).enableTrace).toBe(true); + }); + }); + + describe('processRequest', () => { + const userId = 'user1'; + const sessionId = 'session1'; + const chatHistory: ConversationMessage[] = []; + + it('invokes the flow and returns the decoded response', async () => { + mockClient.send.mockResolvedValueOnce(makeResponseStream('Flow answer') as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + const result = await agent.processRequest('What is 2+2?', userId, sessionId, chatHistory); + + expect(mockClient.send).toHaveBeenCalledTimes(1); + expect(result).toEqual({ + role: ParticipantRole.ASSISTANT, + content: [{ text: 'Flow answer' }], + }); + }); + + it('throws when the response stream is missing', async () => { + mockClient.send.mockResolvedValueOnce({ responseStream: null } as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + await expect(agent.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toThrow(/Error processing request with Bedrock/); + }); + + it('uses a custom flowInputEncoder when provided', async () => { + const encoder = jest.fn().mockReturnValue('encoded input'); + mockClient.send.mockResolvedValueOnce(makeResponseStream('ok') as never); + + const agent = new BedrockFlowsAgent({ ...defaultOptions, flowInputEncoder: encoder }); + await agent.processRequest('hello', userId, sessionId, chatHistory); + + expect(encoder).toHaveBeenCalledWith(agent, 'hello', expect.objectContaining({ userId, sessionId })); + }); + + it('uses a custom flowOutputDecoder when provided', async () => { + const decoded: ConversationMessage = { role: ParticipantRole.ASSISTANT, content: [{ text: 'custom' }] }; + const decoder = jest.fn().mockReturnValue(decoded); + mockClient.send.mockResolvedValueOnce(makeResponseStream('raw') as never); + + const agent = new BedrockFlowsAgent({ ...defaultOptions, flowOutputDecoder: decoder }); + const result = await agent.processRequest('input', userId, sessionId, chatHistory); + + expect(decoder).toHaveBeenCalledWith(agent, 'raw'); + expect(result).toEqual(decoded); + }); + + it('wraps client errors in a descriptive error message', async () => { + mockClient.send.mockRejectedValueOnce(new Error('Network failure') as never); + + const agent = new BedrockFlowsAgent(defaultOptions); + await expect(agent.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toThrow('Error processing request with Bedrock: Network failure'); + }); + }); +}); diff --git a/typescript/tests/agents/ChainAgent.test.ts b/typescript/tests/agents/ChainAgent.test.ts new file mode 100644 index 000000000..d0ac57e82 --- /dev/null +++ b/typescript/tests/agents/ChainAgent.test.ts @@ -0,0 +1,136 @@ +import { ChainAgent, ChainAgentOptions } from '../../src/agents/chainAgent'; +import { Agent, AgentOptions } from '../../src/agents/agent'; +import { ConversationMessage, ParticipantRole } from '../../src/types'; + +// Minimal concrete Agent subclass for testing +class MockAgent extends Agent { + private response: ConversationMessage | AsyncIterable; + + constructor( + options: AgentOptions, + response: ConversationMessage | AsyncIterable + ) { + super(options); + this.response = response; + } + + async processRequest(): Promise> { + return this.response; + } +} + +function makeMessage(text: string): ConversationMessage { + return { role: ParticipantRole.ASSISTANT, content: [{ text }] }; +} + +function makeAgent(name: string, response: ConversationMessage | AsyncIterable): MockAgent { + return new MockAgent({ name, description: `${name} description` }, response); +} + +describe('ChainAgent', () => { + const defaultOptions: ChainAgentOptions = { + name: 'TestChain', + description: 'Test chain agent', + agents: [makeAgent('Agent1', makeMessage('step one output'))], + }; + + describe('constructor', () => { + it('throws when agents array is empty', () => { + expect(() => new ChainAgent({ ...defaultOptions, agents: [] })).toThrow( + 'ChainAgent requires at least one agent in the chain.' + ); + }); + + it('initialises with a single agent', () => { + const chain = new ChainAgent(defaultOptions); + expect(chain.agents).toHaveLength(1); + }); + + it('uses custom defaultOutput when provided', () => { + const chain = new ChainAgent({ ...defaultOptions, defaultOutput: 'custom fallback' }); + expect(chain).toBeDefined(); + }); + }); + + describe('processRequest', () => { + const userId = 'user1'; + const sessionId = 'session1'; + const chatHistory: ConversationMessage[] = []; + + it('returns the output of a single-agent chain', async () => { + const chain = new ChainAgent(defaultOptions); + const result = await chain.processRequest('hello', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.role).toBe(ParticipantRole.ASSISTANT); + expect(result.content[0]).toHaveProperty('text', 'step one output'); + }); + + it('pipes output of each agent as input to the next', async () => { + const spy1 = jest.fn().mockResolvedValue(makeMessage('output of agent 1')); + const spy2 = jest.fn().mockResolvedValue(makeMessage('output of agent 2')); + + const agent1 = new MockAgent({ name: 'A1', description: 'd' }, makeMessage('')); + const agent2 = new MockAgent({ name: 'A2', description: 'd' }, makeMessage('')); + agent1.processRequest = spy1 as any; + agent2.processRequest = spy2 as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [agent1, agent2] }); + await chain.processRequest('initial input', userId, sessionId, chatHistory); + + expect(spy1).toHaveBeenCalledWith('initial input', userId, sessionId, chatHistory, undefined); + expect(spy2).toHaveBeenCalledWith('output of agent 1', userId, sessionId, chatHistory, undefined); + }); + + it('passes additionalParams through to each agent', async () => { + const spy = jest.fn().mockResolvedValue(makeMessage('ok')); + const agent = new MockAgent({ name: 'A', description: 'd' }, makeMessage('')); + agent.processRequest = spy as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [agent] }); + const params = { key: 'val' }; + await chain.processRequest('input', userId, sessionId, chatHistory, params); + + expect(spy).toHaveBeenCalledWith('input', userId, sessionId, chatHistory, params); + }); + + it('returns default response when an intermediate agent returns no text content', async () => { + const emptyAgent = makeAgent('Empty', { role: ParticipantRole.ASSISTANT, content: [] }); + const secondAgent = makeAgent('Second', makeMessage('should not reach')); + + const chain = new ChainAgent({ ...defaultOptions, agents: [emptyAgent, secondAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.content[0]).toHaveProperty('text', 'No output generated from the chain.'); + }); + + it('throws when an agent in the chain errors', async () => { + const errorAgent = new MockAgent({ name: 'ErrorAgent', description: 'd' }, makeMessage('')); + errorAgent.processRequest = jest.fn().mockRejectedValue(new Error('downstream failure')) as any; + + const chain = new ChainAgent({ ...defaultOptions, agents: [errorAgent] }); + await expect(chain.processRequest('input', userId, sessionId, chatHistory)) + .rejects.toMatch(/Error processing request with agent ErrorAgent/); + }); + + it('allows the last agent to return a streaming response', async () => { + async function* streamGen() { yield 'chunk'; } + const streamingAgent = makeAgent('Streamer', streamGen()); + + const chain = new ChainAgent({ ...defaultOptions, agents: [streamingAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory); + + expect(typeof (result as any)[Symbol.asyncIterator]).toBe('function'); + }); + + it('returns default response when a non-last agent returns a streaming response', async () => { + async function* streamGen() { yield 'chunk'; } + const streamingAgent = makeAgent('Streamer', streamGen()); + const nextAgent = makeAgent('Next', makeMessage('next')); + + const chain = new ChainAgent({ ...defaultOptions, agents: [streamingAgent, nextAgent] }); + const result = await chain.processRequest('input', userId, sessionId, chatHistory) as ConversationMessage; + + expect(result.content[0]).toHaveProperty('text', 'No output generated from the chain.'); + }); + }); +});