add retrieve-as-tool mode

This commit is contained in:
Mutasem Aldmour 2024-12-19 14:11:32 +01:00
parent 87e2f2d154
commit 85604cca25
No known key found for this signature in database
GPG key ID: 3DFA8122BB7FD6B8
4 changed files with 109 additions and 12 deletions

View file

@ -228,7 +228,7 @@ export class VectorStorePGVector extends createVectorStoreNode({
testedBy: 'postgresConnectionTest',
},
],
operationModes: ['load', 'insert', 'retrieve'],
operationModes: ['load', 'insert', 'retrieve', 'retrieve-as-tool'],
},
sharedFields,
insertFields,

View file

@ -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,

View file

@ -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 },

View file

@ -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<SupplyData> {
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',
);
}
};