first commit

packages/@n8n/nodes-langchain/nodes/chains/ChainSummarization/V1/ChainSummarizationV1.node.ts

@@ -0,0 +1,262 @@
+import type { Document } from '@langchain/core/documents';
+import type { BaseLanguageModel } from '@langchain/core/language_models/base';
+import { PromptTemplate } from '@langchain/core/prompts';
+import type { SummarizationChainParams } from '@langchain/classic/chains';
+import { loadSummarizationChain } from '@langchain/classic/chains';
+import {
+	NodeConnectionTypes,
+	type INodeTypeBaseDescription,
+	type IExecuteFunctions,
+	type INodeExecutionData,
+	type INodeType,
+	type INodeTypeDescription,
+} from 'n8n-workflow';
+
+import { N8nBinaryLoader, N8nJsonLoader, getTemplateNoticeField } from '@n8n/ai-utilities';
+
+import { REFINE_PROMPT_TEMPLATE, DEFAULT_PROMPT_TEMPLATE } from '../prompt';
+
+export class ChainSummarizationV1 implements INodeType {
+	description: INodeTypeDescription;
+
+	constructor(baseDescription: INodeTypeBaseDescription) {
+		this.description = {
+			...baseDescription,
+			version: 1,
+			defaults: {
+				name: 'Summarization Chain',
+				color: '#909298',
+			},
+
+			inputs: [
+				NodeConnectionTypes.Main,
+				{
+					displayName: 'Model',
+					maxConnections: 1,
+					type: NodeConnectionTypes.AiLanguageModel,
+					required: true,
+				},
+				{
+					displayName: 'Document',
+					maxConnections: 1,
+					type: NodeConnectionTypes.AiDocument,
+					required: true,
+				},
+			],
+			outputs: [NodeConnectionTypes.Main],
+			credentials: [],
+			properties: [
+				getTemplateNoticeField(1951),
+				{
+					displayName: 'Type',
+					name: 'type',
+					type: 'options',
+					description: 'The type of summarization to run',
+					default: 'map_reduce',
+					options: [
+						{
+							name: 'Map Reduce (Recommended)',
+							value: 'map_reduce',
+							description:
+								'Summarize each document (or chunk) individually, then summarize those summaries',
+						},
+						{
+							name: 'Refine',
+							value: 'refine',
+							description:
+								'Summarize the first document (or chunk). Then update that summary based on the next document (or chunk), and repeat.',
+						},
+						{
+							name: 'Stuff',
+							value: 'stuff',
+							description: 'Pass all documents (or chunks) at once. Ideal for small datasets.',
+						},
+					],
+				},
+				{
+					displayName: 'Options',
+					name: 'options',
+					type: 'collection',
+					default: {},
+					placeholder: 'Add Option',
+					options: [
+						{
+							displayName: 'Final Prompt to Combine',
+							name: 'combineMapPrompt',
+							type: 'string',
+							hint: 'The prompt to combine individual summaries',
+							displayOptions: {
+								show: {
+									'/type': ['map_reduce'],
+								},
+							},
+							default: DEFAULT_PROMPT_TEMPLATE,
+							typeOptions: {
+								rows: 6,
+							},
+						},
+						{
+							displayName: 'Individual Summary Prompt',
+							name: 'prompt',
+							type: 'string',
+							default: DEFAULT_PROMPT_TEMPLATE,
+							hint: 'The prompt to summarize an individual document (or chunk)',
+							displayOptions: {
+								show: {
+									'/type': ['map_reduce'],
+								},
+							},
+							typeOptions: {
+								rows: 6,
+							},
+						},
+						{
+							displayName: 'Prompt',
+							name: 'prompt',
+							type: 'string',
+							default: DEFAULT_PROMPT_TEMPLATE,
+							displayOptions: {
+								show: {
+									'/type': ['stuff'],
+								},
+							},
+							typeOptions: {
+								rows: 6,
+							},
+						},
+						{
+							displayName: 'Subsequent (Refine) Prompt',
+							name: 'refinePrompt',
+							type: 'string',
+							displayOptions: {
+								show: {
+									'/type': ['refine'],
+								},
+							},
+							default: REFINE_PROMPT_TEMPLATE,
+							hint: 'The prompt to refine the summary based on the next document (or chunk)',
+							typeOptions: {
+								rows: 6,
+							},
+						},
+						{
+							displayName: 'Initial Prompt',
+							name: 'refineQuestionPrompt',
+							type: 'string',
+							displayOptions: {
+								show: {
+									'/type': ['refine'],
+								},
+							},
+							default: DEFAULT_PROMPT_TEMPLATE,
+							hint: 'The prompt for the first document (or chunk)',
+							typeOptions: {
+								rows: 6,
+							},
+						},
+					],
+				},
+			],
+		};
+	}
+
+	async execute(this: IExecuteFunctions): Promise<INodeExecutionData[][]> {
+		this.logger.debug('Executing Vector Store QA Chain');
+		const type = this.getNodeParameter('type', 0) as 'map_reduce' | 'stuff' | 'refine';
+
+		const model = (await this.getInputConnectionData(
+			NodeConnectionTypes.AiLanguageModel,
+			0,
+		)) as BaseLanguageModel;
+
+		const documentInput = (await this.getInputConnectionData(NodeConnectionTypes.AiDocument, 0)) as
+			| N8nJsonLoader
+			| Array<Document<Record<string, unknown>>>;
+
+		const options = this.getNodeParameter('options', 0, {}) as {
+			prompt?: string;
+			refineQuestionPrompt?: string;
+			refinePrompt?: string;
+			combineMapPrompt?: string;
+		};
+
+		const chainArgs: SummarizationChainParams = {
+			type,
+		};
+
+		// Map reduce prompt override
+		if (type === 'map_reduce') {
+			const mapReduceArgs = chainArgs as SummarizationChainParams & {
+				type: 'map_reduce';
+			};
+			if (options.combineMapPrompt) {
+				mapReduceArgs.combineMapPrompt = new PromptTemplate({
+					template: options.combineMapPrompt,
+					inputVariables: ['text'],
+				});
+			}
+			if (options.prompt) {
+				mapReduceArgs.combinePrompt = new PromptTemplate({
+					template: options.prompt,
+					inputVariables: ['text'],
+				});
+			}
+		}
+
+		// Stuff prompt override
+		if (type === 'stuff') {
+			const stuffArgs = chainArgs as SummarizationChainParams & {
+				type: 'stuff';
+			};
+			if (options.prompt) {
+				stuffArgs.prompt = new PromptTemplate({
+					template: options.prompt,
+					inputVariables: ['text'],
+				});
+			}
+		}
+
+		// Refine prompt override
+		if (type === 'refine') {
+			const refineArgs = chainArgs as SummarizationChainParams & {
+				type: 'refine';
+			};
+
+			if (options.refinePrompt) {
+				refineArgs.refinePrompt = new PromptTemplate({
+					template: options.refinePrompt,
+					inputVariables: ['existing_answer', 'text'],
+				});
+			}
+
+			if (options.refineQuestionPrompt) {
+				refineArgs.questionPrompt = new PromptTemplate({
+					template: options.refineQuestionPrompt,
+					inputVariables: ['text'],
+				});
+			}
+		}
+
+		const chain = loadSummarizationChain(model, chainArgs);
+
+		const items = this.getInputData();
+		const returnData: INodeExecutionData[] = [];
+
+		for (let itemIndex = 0; itemIndex < items.length; itemIndex++) {
+			let processedDocuments: Document[];
+			if (documentInput instanceof N8nJsonLoader || documentInput instanceof N8nBinaryLoader) {
+				processedDocuments = await documentInput.processItem(items[itemIndex], itemIndex);
+			} else {
+				processedDocuments = documentInput;
+			}
+
+			const response = await chain.call({
+				input_documents: processedDocuments,
+			});
+
+			returnData.push({ json: { response } });
+		}
+
+		return [returnData];
+	}
+}