From 85604cca251ffe4f5e7fea48b540afbfbf03b5d5 Mon Sep 17 00:00:00 2001 From: Mutasem Aldmour Date: Thu, 19 Dec 2024 14:11:32 +0100 Subject: [PATCH] add retrieve-as-tool mode --- .../VectorStorePGVector.node.ts | 2 +- .../VectorStorePinecone.node.ts | 2 +- .../VectorStoreSupabase.node.ts | 2 +- .../shared/createVectorStoreNode.ts | 115 ++++++++++++++++-- 4 files changed, 109 insertions(+), 12 deletions(-) diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts index 6d5da1615b..d9d5ee611a 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts @@ -228,7 +228,7 @@ export class VectorStorePGVector extends createVectorStoreNode({ testedBy: 'postgresConnectionTest', }, ], - operationModes: ['load', 'insert', 'retrieve'], + operationModes: ['load', 'insert', 'retrieve', 'retrieve-as-tool'], }, sharedFields, insertFields, diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePinecone/VectorStorePinecone.node.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePinecone/VectorStorePinecone.node.ts index 6e684ebed3..8eb37c69e2 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePinecone/VectorStorePinecone.node.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePinecone/VectorStorePinecone.node.ts @@ -65,7 +65,7 @@ export class VectorStorePinecone extends createVectorStoreNode({ required: true, }, ], - operationModes: ['load', 'insert', 'retrieve', 'update'], + operationModes: ['load', 'insert', 'retrieve', 'update', 'retrieve-as-tool'], }, methods: { listSearch: { pineconeIndexSearch } }, retrieveFields, diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreSupabase/VectorStoreSupabase.node.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreSupabase/VectorStoreSupabase.node.ts index b1b80fea5a..a462ff8cf6 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreSupabase/VectorStoreSupabase.node.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStoreSupabase/VectorStoreSupabase.node.ts @@ -55,7 +55,7 @@ export class VectorStoreSupabase extends createVectorStoreNode({ required: true, }, ], - operationModes: ['load', 'insert', 'retrieve', 'update'], + operationModes: ['load', 'insert', 'retrieve', 'update', 'retrieve-as-tool'], }, methods: { listSearch: { supabaseTableNameSearch }, diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts index 1fff7bef7e..3135a74a08 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode.ts @@ -26,10 +26,16 @@ import { N8nJsonLoader } from '@utils/N8nJsonLoader'; import { getConnectionHintNoticeField } from '@utils/sharedFields'; import { processDocument } from './processDocuments'; +import { DynamicTool } from 'langchain/tools'; -type NodeOperationMode = 'insert' | 'load' | 'retrieve' | 'update'; +type NodeOperationMode = 'insert' | 'load' | 'retrieve' | 'update' | 'retrieve-as-tool'; -const DEFAULT_OPERATION_MODES: NodeOperationMode[] = ['load', 'insert', 'retrieve']; +const DEFAULT_OPERATION_MODES: NodeOperationMode[] = [ + 'load', + 'insert', + 'retrieve', + 'retrieve-as-tool', +]; interface NodeMeta { displayName: string; @@ -100,10 +106,16 @@ function getOperationModeOptions(args: VectorStoreNodeConstructorArgs): INodePro action: 'Add documents to vector store', }, { - name: 'Retrieve Documents (For Agent/Chain)', + name: 'Retrieve Documents (As Vector Store for AI Agent)', value: 'retrieve', - description: 'Retrieve documents from vector store to be used with AI nodes', - action: 'Retrieve documents for AI processing', + description: 'Retrieve documents from vector store to be used as vector store with AI nodes', + action: 'Retrieve documents for AI processing as Vector Store', + }, + { + name: 'Retrieve Documents (As Tool for AI Agent)', + value: 'retrieve-as-tool', + description: 'Retrieve documents from vector store to be used as tool with AI nodes', + action: 'Retrieve documents for AI processing as Tool', }, { name: 'Update Documents', @@ -150,6 +162,10 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => const mode = parameters?.mode; const inputs = [{ displayName: "Embedding", type: "${NodeConnectionType.AiEmbedding}", required: true, maxConnections: 1}] + if (mode === 'retrieve-as-tool') { + return inputs; + } + if (['insert', 'load', 'update'].includes(mode)) { inputs.push({ displayName: "", type: "${NodeConnectionType.Main}"}) } @@ -163,6 +179,11 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => outputs: `={{ ((parameters) => { const mode = parameters?.mode ?? 'retrieve'; + + if (mode === 'retrieve-as-tool') { + return [{ displayName: "Tool", type: "${NodeConnectionType.AiTool}"}] + } + if (mode === 'retrieve') { return [{ displayName: "Vector Store", type: "${NodeConnectionType.AiVectorStore}"}] } @@ -186,6 +207,37 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => }, }, }, + { + displayName: 'Name', + name: 'toolName', + type: 'string', + default: '', + required: true, + description: 'Name of the vector store', + placeholder: 'e.g. company_knowledge_base', + validateType: 'string-alphanumeric', + displayOptions: { + show: { + mode: ['retrieve-as-tool'], + }, + }, + }, + { + displayName: 'Description', + name: 'toolDescription', + type: 'string', + default: '', + required: true, + typeOptions: { rows: 2 }, + description: + 'Explain to the LLM what this tool does, a good, specific description would allow LLMs to produce expected results much more often', + placeholder: `e.g. ${args.meta.description}`, + displayOptions: { + show: { + mode: ['retrieve-as-tool'], + }, + }, + }, ...args.sharedFields, ...transformDescriptionForOperationMode(args.insertFields ?? [], 'insert'), // Prompt and topK are always used for the load operation @@ -211,7 +263,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => description: 'Number of top results to fetch from vector store', displayOptions: { show: { - mode: ['load'], + mode: ['load', 'retrieve-as-tool'], }, }, }, @@ -223,7 +275,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => description: 'Whether or not to include document metadata', displayOptions: { show: { - mode: ['load'], + mode: ['load', 'retrieve-as-tool'], }, }, }, @@ -401,7 +453,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => } async supplyData(this: ISupplyDataFunctions, itemIndex: number): Promise { - const mode = this.getNodeParameter('mode', 0) as 'load' | 'insert' | 'retrieve'; + const mode = this.getNodeParameter('mode', 0) as NodeOperationMode; const filter = getMetadataFiltersValues(this, itemIndex); const embeddings = (await this.getInputConnectionData( NodeConnectionType.AiEmbedding, @@ -415,9 +467,54 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => }; } + if (mode === 'retrieve-as-tool') { + const toolDescription = this.getNodeParameter('toolDescription', itemIndex) as string; + const toolName = this.getNodeParameter('toolName', itemIndex) as string; + const topK = this.getNodeParameter('topK', itemIndex, 4) as number; + const includeDocumentMetadata = this.getNodeParameter( + 'includeDocumentMetadata', + itemIndex, + true, + ) as boolean; + + const vectorStoreTool = new DynamicTool({ + name: toolName, + description: toolDescription, + func: async (input) => { + const vectorStore = await args.getVectorStoreClient( + this, + filter, + embeddings, + itemIndex, + ); + const embeddedPrompt = await embeddings.embedQuery(input); + const documents = await vectorStore.similaritySearchVectorWithScore( + embeddedPrompt, + topK, + filter, + ); + return documents + .map((document) => { + if (includeDocumentMetadata) { + return { type: 'text', text: JSON.stringify(document[0]) }; + } + return { + type: 'text', + text: JSON.stringify({ pageContent: document[0].pageContent }), + }; + }) + .filter((document) => !!document); + }, + }); + + return { + response: logWrapper(vectorStoreTool, this), + }; + } + throw new NodeOperationError( this.getNode(), - 'Only the "retrieve" operation mode is supported to supply data', + 'Only the "retrieve" and "retrieve-as-tool" operation mode is supported to supply data', ); } };