fix(Auto-fixing Output Parser Node): Only run retry chain on parsing …

…errors (#11569)
n8n-io · Nov 6, 2024 · 21b31e4 · 21b31e4
1 parent d960777
commit 21b31e4
Show file tree

Hide file tree

Showing 5 changed files with 203 additions and 60 deletions.
diff --git a/packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/ToolsAgent/execute.ts b/packages/@n8n/nodes-langchain/nodes/agents/Agent/agents/ToolsAgent/execute.ts
@@ -206,10 +206,28 @@ export async function toolsAgentExecute(this: IExecuteFunctions): Promise<INodeE
 		// If the steps are an AgentFinish and the outputParser is defined it must mean that the LLM didn't use `format_final_response` tool so we will try to parse the output manually
 		if (outputParser && typeof steps === 'object' && (steps as AgentFinish).returnValues) {
 			const finalResponse = (steps as AgentFinish).returnValues;
-			const returnValues = (await outputParser.parse(finalResponse as unknown as string)) as Record<
-				string,
-				unknown
-			>;
+			let parserInput: string;
+
+			if (finalResponse instanceof Object) {
+				if ('output' in finalResponse) {
+					try {
+						// If the output is an object, we will try to parse it as JSON
+						// this is because parser expects stringified JSON object like { "output": { .... } }
+						// so we try to parse the output before wrapping it and then stringify it
+						parserInput = JSON.stringify({ output: jsonParse(finalResponse.output) });
+					} catch (error) {
+						// If parsing of the output fails, we will use the raw output
+						parserInput = finalResponse.output;
+					}
+				} else {
+					// If the output is not an object, we will stringify it as it is
+					parserInput = JSON.stringify(finalResponse);
+				}
+			} else {
+				parserInput = finalResponse;
+			}
+
+			const returnValues = (await outputParser.parse(parserInput)) as Record<string, unknown>;
 			return handleParsedStepOutput(returnValues);
 		}
 		return handleAgentFinishOutput(steps);

diff --git a/...nodes-langchain/nodes/output_parser/OutputParserAutofixing/OutputParserAutofixing.node.ts b/...nodes-langchain/nodes/output_parser/OutputParserAutofixing/OutputParserAutofixing.node.ts
@@ -1,12 +1,14 @@
 import type { BaseLanguageModel } from '@langchain/core/language_models/base';
-import { NodeConnectionType } from 'n8n-workflow';
+import { PromptTemplate } from '@langchain/core/prompts';
+import { NodeConnectionType, NodeOperationError } from 'n8n-workflow';
 import type {
 	ISupplyDataFunctions,
 	INodeType,
 	INodeTypeDescription,
 	SupplyData,
 } from 'n8n-workflow';
 
+import { NAIVE_FIX_PROMPT } from './prompt';
 import {
 	N8nOutputFixingParser,
 	type N8nStructuredOutputParser,
@@ -65,6 +67,27 @@ export class OutputParserAutofixing implements INodeType {
 				default: '',
 			},
 			getConnectionHintNoticeField([NodeConnectionType.AiChain, NodeConnectionType.AiAgent]),
+			{
+				displayName: 'Options',
+				name: 'options',
+				type: 'collection',
+				placeholder: 'Add Option',
+				default: {},
+				options: [
+					{
+						displayName: 'Retry Prompt',
+						name: 'prompt',
+						type: 'string',
+						default: NAIVE_FIX_PROMPT,
+						typeOptions: {
+							rows: 10,
+						},
+						hint: 'Should include "{error}", "{instructions}", and "{completion}" placeholders',
+						description:
+							'Prompt template used for fixing the output. Uses placeholders: "{instructions}" for parsing rules, "{completion}" for the failed attempt, and "{error}" for the validation error message.',
+					},
+				],
+			},
 		],
 	};
 
@@ -77,8 +100,20 @@ export class OutputParserAutofixing implements INodeType {
 			NodeConnectionType.AiOutputParser,
 			itemIndex,
 		)) as N8nStructuredOutputParser;
+		const prompt = this.getNodeParameter('options.prompt', itemIndex, NAIVE_FIX_PROMPT) as string;
 
-		const parser = new N8nOutputFixingParser(this, model, outputParser);
+		if (prompt.length === 0 || !prompt.includes('{error}')) {
+			throw new NodeOperationError(
+				this.getNode(),
+				'Auto-fixing parser prompt has to contain {error} placeholder',
+			);
+		}
+		const parser = new N8nOutputFixingParser(
+			this,
+			model,
+			outputParser,
+			PromptTemplate.fromTemplate(prompt),
+		);
 
 		return {
 			response: parser,

diff --git a/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserAutofixing/prompt.ts b/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserAutofixing/prompt.ts
@@ -0,0 +1,16 @@
+export const NAIVE_FIX_PROMPT = `Instructions:
+--------------
+{instructions}
+--------------
+Completion:
+--------------
+{completion}
+--------------
+
+Above, the Completion did not satisfy the constraints given in the Instructions.
+Error:
+--------------
+{error}
+--------------
+
+Please try again. Please only respond with an answer that satisfies the constraints laid out in the Instructions:`;
diff --git a/...chain/nodes/output_parser/OutputParserAutofixing/test/OutputParserAutofixing.node.test.ts b/...chain/nodes/output_parser/OutputParserAutofixing/test/OutputParserAutofixing.node.test.ts
@@ -1,15 +1,19 @@
 /* eslint-disable @typescript-eslint/unbound-method */
 /* eslint-disable @typescript-eslint/no-unsafe-call */
 import type { BaseLanguageModel } from '@langchain/core/language_models/base';
+import { OutputParserException } from '@langchain/core/output_parsers';
 import type { MockProxy } from 'jest-mock-extended';
 import { mock } from 'jest-mock-extended';
 import { normalizeItems } from 'n8n-core';
 import type { IExecuteFunctions, IWorkflowDataProxyData } from 'n8n-workflow';
-import { ApplicationError, NodeConnectionType } from 'n8n-workflow';
+import { ApplicationError, NodeConnectionType, NodeOperationError } from 'n8n-workflow';
 
-import { N8nOutputFixingParser } from '../../../../utils/output_parsers/N8nOutputParser';
-import type { N8nStructuredOutputParser } from '../../../../utils/output_parsers/N8nOutputParser';
+import type {
+	N8nOutputFixingParser,
+	N8nStructuredOutputParser,
+} from '../../../../utils/output_parsers/N8nOutputParser';
 import { OutputParserAutofixing } from '../OutputParserAutofixing.node';
+import { NAIVE_FIX_PROMPT } from '../prompt';
 
 describe('OutputParserAutofixing', () => {
 	let outputParser: OutputParserAutofixing;
@@ -34,6 +38,13 @@ describe('OutputParserAutofixing', () => {
 
 			throw new ApplicationError('Unexpected connection type');
 		});
+		thisArg.getNodeParameter.mockReset();
+		thisArg.getNodeParameter.mockImplementation((parameterName) => {
+			if (parameterName === 'options.prompt') {
+				return NAIVE_FIX_PROMPT;
+			}
+			throw new ApplicationError('Not implemented');
+		});
 	});
 
 	afterEach(() => {
@@ -48,73 +59,132 @@ describe('OutputParserAutofixing', () => {
 		});
 	}
 
-	it('should successfully parse valid output without needing to fix it', async () => {
-		const validOutput = { name: 'Alice', age: 25 };
-
-		mockStructuredOutputParser.parse.mockResolvedValueOnce(validOutput);
+	describe('Configuration', () => {
+		it('should throw error when prompt template does not contain {error} placeholder', async () => {
+			thisArg.getNodeParameter.mockImplementation((parameterName) => {
+				if (parameterName === 'options.prompt') {
+					return 'Invalid prompt without error placeholder';
+				}
+				throw new ApplicationError('Not implemented');
+			});
+
+			await expect(outputParser.supplyData.call(thisArg, 0)).rejects.toThrow(
+				new NodeOperationError(
+					thisArg.getNode(),
+					'Auto-fixing parser prompt has to contain {error} placeholder',
+				),
+			);
+		});
 
-		const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
-			response: N8nOutputFixingParser;
-		};
+		it('should throw error when prompt template is empty', async () => {
+			thisArg.getNodeParameter.mockImplementation((parameterName) => {
+				if (parameterName === 'options.prompt') {
+					return '';
+				}
+				throw new ApplicationError('Not implemented');
+			});
+
+			await expect(outputParser.supplyData.call(thisArg, 0)).rejects.toThrow(
+				new NodeOperationError(
+					thisArg.getNode(),
+					'Auto-fixing parser prompt has to contain {error} placeholder',
+				),
+			);
+		});
 
-		// Ensure the response contains the output-fixing parser
-		expect(response).toBeDefined();
-		expect(response).toBeInstanceOf(N8nOutputFixingParser);
+		it('should use default prompt when none specified', async () => {
+			thisArg.getNodeParameter.mockImplementation((parameterName) => {
+				if (parameterName === 'options.prompt') {
+					return NAIVE_FIX_PROMPT;
+				}
+				throw new ApplicationError('Not implemented');
+			});
 
-		const result = await response.parse('{"name": "Alice", "age": 25}');
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
 
-		// Validate that the parser succeeds without retry
-		expect(result).toEqual(validOutput);
-		expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(1); // Only one call to parse
+			expect(response).toBeDefined();
+		});
 	});
 
-	it('should throw an error when both structured parser and fixing parser fail', async () => {
-		mockStructuredOutputParser.parse
-			.mockRejectedValueOnce(new Error('Invalid JSON')) // First attempt fails
-			.mockRejectedValueOnce(new Error('Fixing attempt failed')); // Second attempt fails
+	describe('Parsing', () => {
+		it('should successfully parse valid output without needing to fix it', async () => {
+			const validOutput = { name: 'Alice', age: 25 };
 
-		const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
-			response: N8nOutputFixingParser;
-		};
+			mockStructuredOutputParser.parse.mockResolvedValueOnce(validOutput);
 
-		response.getRetryChain = getMockedRetryChain('{}');
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
 
-		await expect(response.parse('Invalid JSON string')).rejects.toThrow('Fixing attempt failed');
-		expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(2);
-	});
+			const result = await response.parse('{"name": "Alice", "age": 25}');
 
-	it('should reject on the first attempt and succeed on retry with the parsed content', async () => {
-		const validOutput = { name: 'Bob', age: 28 };
+			expect(result).toEqual(validOutput);
+			expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(1);
+		});
 
-		mockStructuredOutputParser.parse.mockRejectedValueOnce(new Error('Invalid JSON'));
+		it('should not retry on non-OutputParserException errors', async () => {
+			const error = new Error('Some other error');
+			mockStructuredOutputParser.parse.mockRejectedValueOnce(error);
 
-		const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
-			response: N8nOutputFixingParser;
-		};
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
 
-		response.getRetryChain = getMockedRetryChain(JSON.stringify(validOutput));
+			await expect(response.parse('Invalid JSON string')).rejects.toThrow(error);
+			expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(1);
+		});
 
-		mockStructuredOutputParser.parse.mockResolvedValueOnce(validOutput);
+		it('should retry on OutputParserException and succeed', async () => {
+			const validOutput = { name: 'Bob', age: 28 };
 
-		const result = await response.parse('Invalid JSON string');
+			mockStructuredOutputParser.parse
+				.mockRejectedValueOnce(new OutputParserException('Invalid JSON'))
+				.mockResolvedValueOnce(validOutput);
 
-		expect(result).toEqual(validOutput);
-		expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(2); // First fails, second succeeds
-	});
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
 
-	it('should handle non-JSON formatted response from fixing parser', async () => {
-		mockStructuredOutputParser.parse.mockRejectedValueOnce(new Error('Invalid JSON'));
+			response.getRetryChain = getMockedRetryChain(JSON.stringify(validOutput));
 
-		const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
-			response: N8nOutputFixingParser;
-		};
+			const result = await response.parse('Invalid JSON string');
 
-		response.getRetryChain = getMockedRetryChain('This is not JSON');
+			expect(result).toEqual(validOutput);
+			expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(2);
+		});
 
-		mockStructuredOutputParser.parse.mockRejectedValueOnce(new Error('Unexpected token'));
+		it('should handle failed retry attempt', async () => {
+			mockStructuredOutputParser.parse
+				.mockRejectedValueOnce(new OutputParserException('Invalid JSON'))
+				.mockRejectedValueOnce(new Error('Still invalid JSON'));
 
-		// Expect the structured parser to throw an error on invalid JSON from retry
-		await expect(response.parse('Invalid JSON string')).rejects.toThrow('Unexpected token');
-		expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(2); // First fails, second tries and fails
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
+
+			response.getRetryChain = getMockedRetryChain('Still not valid JSON');
+
+			await expect(response.parse('Invalid JSON string')).rejects.toThrow('Still invalid JSON');
+			expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(2);
+		});
+
+		it('should throw non-OutputParserException errors immediately without retry', async () => {
+			const customError = new Error('Database connection error');
+			const retryChainSpy = jest.fn();
+
+			mockStructuredOutputParser.parse.mockRejectedValueOnce(customError);
+
+			const { response } = (await outputParser.supplyData.call(thisArg, 0)) as {
+				response: N8nOutputFixingParser;
+			};
+
+			response.getRetryChain = retryChainSpy;
+
+			await expect(response.parse('Some input')).rejects.toThrow('Database connection error');
+			expect(mockStructuredOutputParser.parse).toHaveBeenCalledTimes(1);
+			expect(retryChainSpy).not.toHaveBeenCalled();
+		});
 	});
 });
diff --git a/packages/@n8n/nodes-langchain/utils/output_parsers/N8nOutputFixingParser.ts b/packages/@n8n/nodes-langchain/utils/output_parsers/N8nOutputFixingParser.ts
@@ -1,12 +1,12 @@
 import type { Callbacks } from '@langchain/core/callbacks/manager';
 import type { BaseLanguageModel } from '@langchain/core/language_models/base';
 import type { AIMessage } from '@langchain/core/messages';
-import { BaseOutputParser } from '@langchain/core/output_parsers';
+import { BaseOutputParser, OutputParserException } from '@langchain/core/output_parsers';
+import type { PromptTemplate } from '@langchain/core/prompts';
 import type { ISupplyDataFunctions } from 'n8n-workflow';
 import { NodeConnectionType } from 'n8n-workflow';
 
 import type { N8nStructuredOutputParser } from './N8nStructuredOutputParser';
-import { NAIVE_FIX_PROMPT } from './prompt';
 import { logAiEvent } from '../helpers';
 
 export class N8nOutputFixingParser extends BaseOutputParser {
@@ -16,12 +16,13 @@ export class N8nOutputFixingParser extends BaseOutputParser {
 		private context: ISupplyDataFunctions,
 		private model: BaseLanguageModel,
 		private outputParser: N8nStructuredOutputParser,
+		private fixPromptTemplate: PromptTemplate,
 	) {
 		super();
 	}
 
 	getRetryChain() {
-		return NAIVE_FIX_PROMPT.pipe(this.model);
+		return this.fixPromptTemplate.pipe(this.model);
 	}
 
 	/**
@@ -47,11 +48,14 @@ export class N8nOutputFixingParser extends BaseOutputParser {
 
 			return response;
 		} catch (error) {
+			if (!(error instanceof OutputParserException)) {
+				throw error;
+			}
 			try {
 				// Second attempt: use retry chain to fix the output
 				const result = (await this.getRetryChain().invoke({
 					completion,
-					error,
+					error: error.message,
 					instructions: this.getFormatInstructions(),
 				})) as AIMessage;