mirror of
https://github.com/Abdulazizzn/n8n-enterprise-unlocked.git
synced 2025-12-16 17:46:45 +00:00
fix(AI Agent Node): Respect context window length in streaming mode (#19567)
This commit is contained in:
@@ -265,7 +265,13 @@ export async function toolsAgentExecute(
|
|||||||
isStreamingAvailable &&
|
isStreamingAvailable &&
|
||||||
this.getNode().typeVersion >= 2.1
|
this.getNode().typeVersion >= 2.1
|
||||||
) {
|
) {
|
||||||
const chatHistory = await memory?.chatHistory.getMessages();
|
// Get chat history respecting the context window length configured in memory
|
||||||
|
let chatHistory;
|
||||||
|
if (memory) {
|
||||||
|
// Load memory variables to respect context window length
|
||||||
|
const memoryVariables = await memory.loadMemoryVariables({});
|
||||||
|
chatHistory = memoryVariables['chat_history'];
|
||||||
|
}
|
||||||
const eventStream = executor.streamEvents(
|
const eventStream = executor.streamEvents(
|
||||||
{
|
{
|
||||||
...invokeParams,
|
...invokeParams,
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import type { ISupplyDataFunctions, IExecuteFunctions, INode } from 'n8n-workflo
|
|||||||
|
|
||||||
import * as helpers from '../../../../../utils/helpers';
|
import * as helpers from '../../../../../utils/helpers';
|
||||||
import * as outputParserModule from '../../../../../utils/output_parsers/N8nOutputParser';
|
import * as outputParserModule from '../../../../../utils/output_parsers/N8nOutputParser';
|
||||||
|
import * as commonModule from '../../agents/ToolsAgent/common';
|
||||||
import { toolsAgentExecute } from '../../agents/ToolsAgent/V2/execute';
|
import { toolsAgentExecute } from '../../agents/ToolsAgent/V2/execute';
|
||||||
|
|
||||||
jest.mock('../../../../../utils/output_parsers/N8nOutputParser', () => ({
|
jest.mock('../../../../../utils/output_parsers/N8nOutputParser', () => ({
|
||||||
@@ -13,6 +14,11 @@ jest.mock('../../../../../utils/output_parsers/N8nOutputParser', () => ({
|
|||||||
N8nStructuredOutputParser: jest.fn(),
|
N8nStructuredOutputParser: jest.fn(),
|
||||||
}));
|
}));
|
||||||
|
|
||||||
|
jest.mock('../../agents/ToolsAgent/common', () => ({
|
||||||
|
...jest.requireActual('../../agents/ToolsAgent/common'),
|
||||||
|
getOptionalMemory: jest.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
const mockHelpers = mock<IExecuteFunctions['helpers']>();
|
const mockHelpers = mock<IExecuteFunctions['helpers']>();
|
||||||
const mockContext = mock<IExecuteFunctions>({ helpers: mockHelpers });
|
const mockContext = mock<IExecuteFunctions>({ helpers: mockHelpers });
|
||||||
|
|
||||||
@@ -620,6 +626,65 @@ describe('toolsAgentExecute', () => {
|
|||||||
expect(result[0][0].json.output).toBe('Regular response');
|
expect(result[0][0].json.output).toBe('Regular response');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('should respect context window length from memory in streaming mode', async () => {
|
||||||
|
const mockMemory = {
|
||||||
|
loadMemoryVariables: jest.fn().mockResolvedValue({
|
||||||
|
chat_history: [
|
||||||
|
{ role: 'human', content: 'Message 1' },
|
||||||
|
{ role: 'ai', content: 'Response 1' },
|
||||||
|
],
|
||||||
|
}),
|
||||||
|
chatHistory: {
|
||||||
|
getMessages: jest.fn().mockResolvedValue([
|
||||||
|
{ role: 'human', content: 'Message 1' },
|
||||||
|
{ role: 'ai', content: 'Response 1' },
|
||||||
|
{ role: 'human', content: 'Message 2' },
|
||||||
|
{ role: 'ai', content: 'Response 2' },
|
||||||
|
]),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
jest.spyOn(commonModule, 'getOptionalMemory').mockResolvedValue(mockMemory as any);
|
||||||
|
|
||||||
|
jest.spyOn(helpers, 'getConnectedTools').mockResolvedValue([mock<Tool>()]);
|
||||||
|
jest.spyOn(outputParserModule, 'getOptionalOutputParser').mockResolvedValue(undefined);
|
||||||
|
mockContext.isStreaming.mockReturnValue(true);
|
||||||
|
|
||||||
|
const mockStreamEvents = async function* () {
|
||||||
|
yield {
|
||||||
|
event: 'on_chat_model_stream',
|
||||||
|
data: {
|
||||||
|
chunk: {
|
||||||
|
content: 'Response',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
};
|
||||||
|
|
||||||
|
const mockExecutor = {
|
||||||
|
streamEvents: jest.fn().mockReturnValue(mockStreamEvents()),
|
||||||
|
};
|
||||||
|
|
||||||
|
jest.spyOn(AgentExecutor, 'fromAgentAndTools').mockReturnValue(mockExecutor as any);
|
||||||
|
|
||||||
|
await toolsAgentExecute.call(mockContext);
|
||||||
|
|
||||||
|
// Verify that memory.loadMemoryVariables was called instead of chatHistory.getMessages
|
||||||
|
expect(mockMemory.loadMemoryVariables).toHaveBeenCalledWith({});
|
||||||
|
expect(mockMemory.chatHistory.getMessages).not.toHaveBeenCalled();
|
||||||
|
|
||||||
|
// Verify that streamEvents was called with the filtered chat history from loadMemoryVariables
|
||||||
|
expect(mockExecutor.streamEvents).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({
|
||||||
|
chat_history: [
|
||||||
|
{ role: 'human', content: 'Message 1' },
|
||||||
|
{ role: 'ai', content: 'Response 1' },
|
||||||
|
],
|
||||||
|
}),
|
||||||
|
expect.any(Object),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
it('should handle mixed message content types in streaming', async () => {
|
it('should handle mixed message content types in streaming', async () => {
|
||||||
jest.spyOn(helpers, 'getConnectedTools').mockResolvedValue([mock<Tool>()]);
|
jest.spyOn(helpers, 'getConnectedTools').mockResolvedValue([mock<Tool>()]);
|
||||||
jest.spyOn(outputParserModule, 'getOptionalOutputParser').mockResolvedValue(undefined);
|
jest.spyOn(outputParserModule, 'getOptionalOutputParser').mockResolvedValue(undefined);
|
||||||
|
|||||||
Reference in New Issue
Block a user