From 7d82dc1ea8c86c2a3575c3629b29af3cfb709f44 Mon Sep 17 00:00:00 2001
From: oleg <me@olegivaniv.com>
Date: Fri, 1 Mar 2024 08:41:45 +0100
Subject: [PATCH] fix: Simplify Structured Output Parser wrapping and fix
 auto-fixing output parser (#8778)

Signed-off-by: Oleg Ivaniv <me@olegivaniv.com>
---
 .../OutputParserStructured.node.ts            | 70 ++++++++++++-------
 .../@n8n/nodes-langchain/utils/logWrapper.ts  |  6 +-
 2 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserStructured/OutputParserStructured.node.ts b/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserStructured/OutputParserStructured.node.ts
index 6ee592ef39..e4fa7abfc8 100644
--- a/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserStructured/OutputParserStructured.node.ts
+++ b/packages/@n8n/nodes-langchain/nodes/output_parser/OutputParserStructured/OutputParserStructured.node.ts
@@ -41,6 +41,7 @@ class N8nStructuredOutputParser<T extends z.ZodTypeAny> extends StructuredOutput
 
 	static fromZedJsonSchema(
 		schema: JSONSchema7,
+		nodeVersion: number,
 	): StructuredOutputParser<z.ZodType<object, z.ZodTypeDef, object>> {
 		// Make sure to remove the description from root schema
 		const { description, ...restOfSchema } = schema;
@@ -51,30 +52,37 @@ class N8nStructuredOutputParser<T extends z.ZodTypeAny> extends StructuredOutput
 		// eslint-disable-next-line @typescript-eslint/no-implied-eval
 		const itemSchema = new Function('z', `return (${zodSchemaString})`)(z) as z.ZodSchema<object>;
 
-		const returnSchema = z.object({
-			[STRUCTURED_OUTPUT_KEY]: z
-				.object({
-					[STRUCTURED_OUTPUT_OBJECT_KEY]: itemSchema.optional(),
-					[STRUCTURED_OUTPUT_ARRAY_KEY]: z.array(itemSchema).optional(),
-				})
-				.describe(
-					`Wrapper around the output data. It can only contain ${STRUCTURED_OUTPUT_OBJECT_KEY} or ${STRUCTURED_OUTPUT_ARRAY_KEY} but never both.`,
-				)
-				.refine(
-					(data) => {
-						// Validate that one and only one of the properties exists
-						return (
-							Boolean(data[STRUCTURED_OUTPUT_OBJECT_KEY]) !==
-							Boolean(data[STRUCTURED_OUTPUT_ARRAY_KEY])
-						);
-					},
-					{
-						message:
-							'One and only one of __structured__output__object and __structured__output__array should be present.',
-						path: [STRUCTURED_OUTPUT_KEY],
-					},
-				),
-		});
+		let returnSchema: z.ZodSchema<object>;
+		if (nodeVersion === 1) {
+			returnSchema = z.object({
+				[STRUCTURED_OUTPUT_KEY]: z
+					.object({
+						[STRUCTURED_OUTPUT_OBJECT_KEY]: itemSchema.optional(),
+						[STRUCTURED_OUTPUT_ARRAY_KEY]: z.array(itemSchema).optional(),
+					})
+					.describe(
+						`Wrapper around the output data. It can only contain ${STRUCTURED_OUTPUT_OBJECT_KEY} or ${STRUCTURED_OUTPUT_ARRAY_KEY} but never both.`,
+					)
+					.refine(
+						(data) => {
+							// Validate that one and only one of the properties exists
+							return (
+								Boolean(data[STRUCTURED_OUTPUT_OBJECT_KEY]) !==
+								Boolean(data[STRUCTURED_OUTPUT_ARRAY_KEY])
+							);
+						},
+						{
+							message:
+								'One and only one of __structured__output__object and __structured__output__array should be present.',
+							path: [STRUCTURED_OUTPUT_KEY],
+						},
+					),
+			});
+		} else {
+			returnSchema = z.object({
+				output: itemSchema.optional(),
+			});
+		}
 
 		return N8nStructuredOutputParser.fromZodSchema(returnSchema);
 	}
@@ -85,7 +93,8 @@ export class OutputParserStructured implements INodeType {
 		name: 'outputParserStructured',
 		icon: 'fa:code',
 		group: ['transform'],
-		version: 1,
+		version: [1, 1.1],
+		defaultVersion: 1.1,
 		description: 'Return data in a defined JSON format',
 		defaults: {
 			name: 'Structured Output Parser',
@@ -152,11 +161,20 @@ export class OutputParserStructured implements INodeType {
 		let itemSchema: JSONSchema7;
 		try {
 			itemSchema = jsonParse<JSONSchema7>(schema);
+
+			// If the type is not defined, we assume it's an object
+			if (itemSchema.type === undefined) {
+				itemSchema = {
+					type: 'object',
+					properties: itemSchema.properties || (itemSchema as { [key: string]: JSONSchema7 }),
+				};
+			}
 		} catch (error) {
 			throw new NodeOperationError(this.getNode(), 'Error during parsing of JSON Schema.');
 		}
 
-		const parser = N8nStructuredOutputParser.fromZedJsonSchema(itemSchema);
+		const nodeVersion = this.getNode().typeVersion;
+		const parser = N8nStructuredOutputParser.fromZedJsonSchema(itemSchema, nodeVersion);
 
 		return {
 			response: logWrapper(parser, this),
diff --git a/packages/@n8n/nodes-langchain/utils/logWrapper.ts b/packages/@n8n/nodes-langchain/utils/logWrapper.ts
index 2b358f7d68..70e30b5ce9 100644
--- a/packages/@n8n/nodes-langchain/utils/logWrapper.ts
+++ b/packages/@n8n/nodes-langchain/utils/logWrapper.ts
@@ -18,7 +18,7 @@ import { BaseChatMemory } from 'langchain/memory';
 import type { MemoryVariables } from 'langchain/dist/memory/base';
 import { BaseRetriever } from 'langchain/schema/retriever';
 import type { FormatInstructionsOptions } from 'langchain/schema/output_parser';
-import { BaseOutputParser } from 'langchain/schema/output_parser';
+import { BaseOutputParser, OutputParserException } from 'langchain/schema/output_parser';
 import { isObject } from 'lodash';
 import { N8nJsonLoader } from './N8nJsonLoader';
 import { N8nBinaryLoader } from './N8nBinaryLoader';
@@ -44,6 +44,10 @@ export async function callMethodAsync<T>(
 	try {
 		return await parameters.method.call(this, ...parameters.arguments);
 	} catch (e) {
+		// Langchain checks for OutputParserException to run retry chain
+		// for auto-fixing the output so skip wrapping in this case
+		if (e instanceof OutputParserException) throw e;
+
 		// Propagate errors from sub-nodes
 		if (e.functionality === 'configuration-node') throw e;
 		const connectedNode = parameters.executeFunctions.getNode();