From df2bc84d2b3830d31319c108f1b01256de95e774 Mon Sep 17 00:00:00 2001
From: oleg <me@olegivaniv.com>
Date: Thu, 27 Jun 2024 13:21:03 +0200
Subject: [PATCH] feat(Vector Store Tool Node): Add Vector Store Tool (#9865)

Signed-off-by: Oleg Ivaniv <me@olegivaniv.com>
---
 .../ToolVectorStore/ToolVectorStore.node.ts   | 113 ++++++++++++++++++
 packages/@n8n/nodes-langchain/package.json    |   1 +
 2 files changed, 114 insertions(+)
 create mode 100644 packages/@n8n/nodes-langchain/nodes/tools/ToolVectorStore/ToolVectorStore.node.ts

diff --git a/packages/@n8n/nodes-langchain/nodes/tools/ToolVectorStore/ToolVectorStore.node.ts b/packages/@n8n/nodes-langchain/nodes/tools/ToolVectorStore/ToolVectorStore.node.ts
new file mode 100644
index 0000000000..abdab09053
--- /dev/null
+++ b/packages/@n8n/nodes-langchain/nodes/tools/ToolVectorStore/ToolVectorStore.node.ts
@@ -0,0 +1,113 @@
+import type { IExecuteFunctions, INodeType, INodeTypeDescription, SupplyData } from 'n8n-workflow';
+import { NodeConnectionType } from 'n8n-workflow';
+
+import { VectorStoreQATool } from 'langchain/tools';
+import type { VectorStore } from '@langchain/core/vectorstores';
+import type { BaseLanguageModel } from '@langchain/core/language_models/base';
+import { VectorDBQAChain } from 'langchain/chains';
+import { getConnectionHintNoticeField } from '../../../utils/sharedFields';
+import { logWrapper } from '../../../utils/logWrapper';
+
+export class ToolVectorStore implements INodeType {
+	description: INodeTypeDescription = {
+		displayName: 'Vector Store Tool',
+		name: 'toolVectorStore',
+		icon: 'fa:database',
+		group: ['transform'],
+		version: [1],
+		description: 'Retrieve context from vector store',
+		defaults: {
+			name: 'Vector Store Tool',
+		},
+		codex: {
+			categories: ['AI'],
+			subcategories: {
+				AI: ['Tools'],
+			},
+			resources: {
+				primaryDocumentation: [
+					{
+						url: 'https://docs.n8n.io/integrations/builtin/cluster-nodes/sub-nodes/n8n-nodes-langchain.toolvectorstore/',
+					},
+				],
+			},
+		},
+		// eslint-disable-next-line n8n-nodes-base/node-class-description-inputs-wrong-regular-node
+		inputs: [
+			{
+				displayName: 'Vector Store',
+				maxConnections: 1,
+				type: NodeConnectionType.AiVectorStore,
+				required: true,
+			},
+			{
+				displayName: 'Model',
+				maxConnections: 1,
+				type: NodeConnectionType.AiLanguageModel,
+				required: true,
+			},
+		],
+		// eslint-disable-next-line n8n-nodes-base/node-class-description-outputs-wrong
+		outputs: [NodeConnectionType.AiTool],
+		outputNames: ['Tool'],
+		properties: [
+			getConnectionHintNoticeField([NodeConnectionType.AiAgent]),
+			{
+				displayName: 'Name',
+				name: 'name',
+				type: 'string',
+				default: '',
+				placeholder: 'e.g. state_of_union_address',
+				validateType: 'string-alphanumeric',
+				description: 'Name of the vector store',
+			},
+			{
+				displayName: 'Description',
+				name: 'description',
+				type: 'string',
+				default: '',
+				placeholder: 'The most recent state of the Union address',
+				typeOptions: {
+					rows: 3,
+				},
+			},
+			{
+				displayName: 'Limit',
+				name: 'topK',
+				type: 'number',
+				default: 4,
+				description: 'The maximum number of results to return',
+			},
+		],
+	};
+
+	async supplyData(this: IExecuteFunctions, itemIndex: number): Promise<SupplyData> {
+		const name = this.getNodeParameter('name', itemIndex) as string;
+		const toolDescription = this.getNodeParameter('description', itemIndex) as string;
+		const topK = this.getNodeParameter('topK', itemIndex, 4) as number;
+
+		const vectorStore = (await this.getInputConnectionData(
+			NodeConnectionType.AiVectorStore,
+			itemIndex,
+		)) as VectorStore;
+
+		const llm = (await this.getInputConnectionData(
+			NodeConnectionType.AiLanguageModel,
+			0,
+		)) as BaseLanguageModel;
+
+		const description = VectorStoreQATool.getDescription(name, toolDescription);
+		const vectorStoreTool = new VectorStoreQATool(name, description, {
+			llm,
+			vectorStore,
+		});
+
+		vectorStoreTool.chain = VectorDBQAChain.fromLLM(llm, vectorStore, {
+			k: topK,
+		});
+
+		return {
+			response: logWrapper(vectorStoreTool, this),
+		};
+	}
+}
diff --git a/packages/@n8n/nodes-langchain/package.json b/packages/@n8n/nodes-langchain/package.json
index c8d486b125..9e0f07c094 100644
--- a/packages/@n8n/nodes-langchain/package.json
+++ b/packages/@n8n/nodes-langchain/package.json
@@ -100,6 +100,7 @@
       "dist/nodes/tools/ToolCode/ToolCode.node.js",
       "dist/nodes/tools/ToolHttpRequest/ToolHttpRequest.node.js",
       "dist/nodes/tools/ToolSerpApi/ToolSerpApi.node.js",
+      "dist/nodes/tools/ToolVectorStore/ToolVectorStore.node.js",
       "dist/nodes/tools/ToolWikipedia/ToolWikipedia.node.js",
       "dist/nodes/tools/ToolWolframAlpha/ToolWolframAlpha.node.js",
       "dist/nodes/tools/ToolWorkflow/ToolWorkflow.node.js",