From d857176abca4b85b4e070e550da22e864b3641bb Mon Sep 17 00:00:00 2001
From: Eugene Molodkin <eugene@n8n.io>
Date: Tue, 3 Sep 2024 16:53:16 +0200
Subject: [PATCH] wip: add auto-fixing parser

---
 .../chains/TextClassifier/TextClassifier.node.ts | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/packages/@n8n/nodes-langchain/nodes/chains/TextClassifier/TextClassifier.node.ts b/packages/@n8n/nodes-langchain/nodes/chains/TextClassifier/TextClassifier.node.ts
index 8765c91b1f666..ab6d4dedd3a14 100644
--- a/packages/@n8n/nodes-langchain/nodes/chains/TextClassifier/TextClassifier.node.ts
+++ b/packages/@n8n/nodes-langchain/nodes/chains/TextClassifier/TextClassifier.node.ts
@@ -12,7 +12,7 @@ import { NodeConnectionType } from 'n8n-workflow';
 import type { BaseLanguageModel } from '@langchain/core/language_models/base';
 import { HumanMessage } from '@langchain/core/messages';
 import { SystemMessagePromptTemplate, ChatPromptTemplate } from '@langchain/core/prompts';
-import { StructuredOutputParser } from 'langchain/output_parsers';
+import { OutputFixingParser, StructuredOutputParser } from 'langchain/output_parsers';
 import { z } from 'zod';
 import { getTracingConfig } from '../../../utils/tracing';
 
@@ -151,6 +151,13 @@ export class TextClassifier implements INodeType {
 							rows: 6,
 						},
 					},
+					{
+						displayName: 'Enable Auto-Fixing',
+						name: 'enableAutoFixing',
+						type: 'boolean',
+						default: true,
+						description: 'Whether to enable auto-fixing for the output parser',
+					},
 				],
 			},
 		],
@@ -173,6 +180,7 @@ export class TextClassifier implements INodeType {
 			multiClass: boolean;
 			fallback?: string;
 			systemPromptTemplate?: string;
+			enableAutoFixing: boolean;
 		};
 		const multiClass = options?.multiClass ?? false;
 		const fallback = options?.fallback ?? 'discard';
@@ -192,7 +200,11 @@ export class TextClassifier implements INodeType {
 			]);
 		const schema = z.object(Object.fromEntries(schemaEntries));
 
-		const parser = StructuredOutputParser.fromZodSchema(schema);
+		const structuredParser = StructuredOutputParser.fromZodSchema(schema);
+
+		const parser = options.enableAutoFixing
+			? OutputFixingParser.fromLLM(llm, structuredParser)
+			: structuredParser;
 
 		const multiClassPrompt = multiClass
 			? 'Categories are not mutually exclusive, and multiple can be true'