feat: Support vector stores as tools

This commit is contained in:
Mutasem Aldmour 2024-12-10 20:10:54 +01:00
parent 956b11a560
commit 819a397735
No known key found for this signature in database
GPG key ID: 3DFA8122BB7FD6B8
3 changed files with 79 additions and 12 deletions

View file

@ -25,6 +25,7 @@ import { logWrapper } from '../../../utils/logWrapper';
import type { N8nBinaryLoader } from '../../../utils/N8nBinaryLoader';
import { N8nJsonLoader } from '../../../utils/N8nJsonLoader';
import { getConnectionHintNoticeField } from '../../../utils/sharedFields';
import { DynamicTool } from 'langchain/tools';
type NodeOperationMode = 'insert' | 'load' | 'retrieve' | 'update';
@ -124,11 +125,12 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
name: args.meta.name,
description: args.meta.description,
icon: args.meta.icon,
group: ['transform'],
group: ['transform', 'vector-store'],
version: 1,
defaults: {
name: args.meta.displayName,
},
usableAsTool: true,
codex: {
categories: ['AI'],
subcategories: {
@ -210,7 +212,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
description: 'Number of top results to fetch from vector store',
displayOptions: {
show: {
mode: ['load'],
mode: ['load', 'retrieve'],
},
},
},
@ -377,7 +379,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
throw new NodeOperationError(
this.getNode(),
'Only the "load" and "insert" operation modes are supported with execute',
'Only the "load", "update" and "insert" operation modes are supported with execute',
);
}
@ -389,16 +391,59 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
0,
)) as Embeddings;
if (mode === 'retrieve') {
const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex);
if (mode !== 'retrieve') {
throw new NodeOperationError(
this.getNode(),
'Only the "retrieve" operation mode is supported to supply data',
);
}
// "Tool" is appended automatically to node type name
// because of usableAsTool flag set above
const isTool = this.getNode().type.endsWith('Tool');
if (isTool) {
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 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) => {
// Tools can only return a string or array of objects with type text
// todo return concatenated strings instead?
return { type: 'text', text: document[0].pageContent };
// todo with metadata?
// return { type: 'text', text: JSON.stringify(document[0]) };
})
.filter((document) => !!document);
},
});
return {
response: logWrapper(vectorStore, this),
response: logWrapper(vectorStoreTool, this),
};
}
throw new NodeOperationError(
this.getNode(),
'Only the "retrieve" operation mode is supported to supply data',
);
const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex);
return {
response: logWrapper(vectorStore, this),
};
}
};

View file

@ -46,7 +46,9 @@ export class NodeTypes implements INodeTypes {
getByNameAndVersion(nodeType: string, version?: number): INodeType {
const origType = nodeType;
const toolRequested = nodeType.startsWith('n8n-nodes-base') && nodeType.endsWith('Tool');
const toolRequested =
(nodeType.startsWith('n8n-nodes-base') || nodeType.startsWith('@n8n/n8n-nodes-langchain')) &&
nodeType.endsWith('Tool');
// Make sure the nodeType to actually get from disk is the un-wrapped type
if (toolRequested) {
nodeType = nodeType.replace(/Tool$/, '');

View file

@ -365,8 +365,13 @@ export function convertNodeToAiTool<
}
if (isFullDescription(item.description)) {
const isVectorStore = item.description.group.includes('vector-store');
item.description.name += 'Tool';
item.description.inputs = [];
if (!isVectorStore) {
item.description.inputs = [];
}
item.description.outputs = [NodeConnectionType.AiTool];
item.description.displayName += ' Tool';
delete item.description.usableAsTool;
@ -417,6 +422,21 @@ export function convertNodeToAiTool<
item.description.properties.unshift(descProp);
if (isVectorStore) {
const nameProp: INodeProperties = {
displayName: 'Name',
name: 'toolName',
type: 'string',
default: '',
required: true,
description: 'Name of the vector store',
placeholder: 'e.g. company_knowledge_base',
validateType: 'string-alphanumeric',
};
item.description.properties.unshift(nameProp);
}
// If node has resource or operation we can determine pre-populate tool description based on it
// so we add the descriptionType property as the first property
if (hasResource || hasOperation) {