N8n Custom AI Agent Tool + Azure AI Search

Hey everyone. I am trying to build an AI Agent tool for Azure AI Search as a vector store. I took this GitHub and the existing Milvus Vector Store as the inspiration and trying to add ‘‘retrieve-as-tool’‘ capability.

I had then added the credentials, nodes and the utils alongiwith package.json to a new docker image and ran a container using it.

The node shows up on canvas but what ever I do it isn’t invoked. Can anyone please guide me on this- this is my first build with n8n.

createVectorStoreNode.ts:

// from @n8n/n8n-nodes-langchain:1.48.0
/* eslint-disable n8n-nodes-base/node-filename-against-convention */
/* eslint-disable n8n-nodes-base/node-dirname-against-convention */
import type { VectorStore } from '@langchain/core/vectorstores';
import { SearchClient } from "@azure/search-documents";
import type { SearchResult } from "@azure/search-documents";
import type { PagedAsyncIterableIterator } from "@azure/core-paging";
import { NodeConnectionType, NodeOperationError } from 'n8n-workflow';
import type {
	INodeCredentialDescription,
	INodeProperties,
	INodeExecutionData,
	IExecuteFunctions,
	INodeTypeDescription,
	SupplyData,
	ISupplyDataFunctions,
	INodeType,
	ILoadOptionsFunctions,
	INodeListSearchResult,
	Icon,
	INode
} from 'n8n-workflow';
import type { Embeddings } from '@langchain/core/embeddings';
import { Document } from '@langchain/core/documents';
import { DynamicTool } from 'langchain/tools';
import { logWrapper } from '../../../utils/logWrapper';
import type { N8nJsonLoader } from '../../../utils/N8nJsonLoader';
import type { N8nBinaryLoader } from '../../../utils/N8nBinaryLoader';
import { getMetadataFiltersValues, logAiEvent } from '../../../utils/helpers';
import { getConnectionHintNoticeField } from '../../../utils/sharedFields';
import { processDocument } from './processDocuments';

interface SearchDocument {
	source?: string;
	content?: string;
	category?: string;
	sourcePage?: string;
	url?: string;
	classification?: string;
	id?: string;
	title?: string;
	documentNumber?: string;
	contentVector?: number[];
	additionalContentVector?: number[];
}

interface NodeMeta {
	displayName: string;
	name: string;
	description: string;
	docsUrl: string;
	icon: Icon;
	credentials?: INodeCredentialDescription[];
}
interface VectorStoreNodeConstructorArgs {
	meta: NodeMeta;
	methods?: {
		listSearch?: {
			[key: string]: (
				this: ILoadOptionsFunctions,
				filter?: string,
				paginationToken?: string,
			) => Promise<INodeListSearchResult>;
		};
	};
	sharedFields: INodeProperties[];
	insertFields?: INodeProperties[];
	loadFields?: INodeProperties[];
	retrieveFields?: INodeProperties[];
	populateVectorStore: (
		context: IExecuteFunctions | ISupplyDataFunctions,
		embeddings: Embeddings,
		documents: Array<Document<Record<string, unknown>>>,
		itemIndex: number,
	) => Promise<void>;
	getVectorStoreClient: (
		context: IExecuteFunctions | ISupplyDataFunctions,
		filter: Record<string, never> | undefined,
		embeddings: Embeddings,
		itemIndex: number,
	) => Promise<VectorStore>;
	getSearchClient: (
		context: IExecuteFunctions | ISupplyDataFunctions,
		itemIndex: number,
	) => Promise<SearchClient<SearchDocument>>
	releaseVectorStoreClient?: (vectorStore: VectorStore) => void;
}

function transformDescriptionForOperationMode(
	fields: INodeProperties[],
	mode: 'insert' | 'load' | 'retrieve' | 'retrieve-as-tool',
) {
	return fields.map((field) => ({
		...field,
		displayOptions: { show: { mode: [mode] } },
	}));
}

function nodeNameToToolName(nodeOrName: INode | string): string {
	const rawName =
		typeof nodeOrName === 'string'
			? nodeOrName
			: nodeOrName?.name;

	const name = rawName && rawName.trim().length > 0 ? rawName : 'search_knowledgebase';
	return name.replace(/[^a-zA-Z0-9_-]+/g, '_');
}

export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
	class VectorStoreNodeType implements INodeType {
		description: INodeTypeDescription = {
			displayName: args.meta.displayName,
			name: args.meta.name,
			description: args.meta.description,
			icon: args.meta.icon,
			group: ['transform'],
			version: 1,
			defaults: {
				name: args.meta.displayName,
			},
			codex: {
				categories: ['AI'],
				subcategories: {
					AI: ['Vector Stores', 'Tools', 'Root Nodes'],
					'Vector Stores': ['Other Vector Stores'],
					Tools: ['Other Tools'],
				},
				resources: {
					primaryDocumentation: [
						{
							url: args.meta.docsUrl,
						},
					],
				},
			},
			credentials: args.meta.credentials,
			// eslint-disable-next-line n8n-nodes-base/node-class-description-inputs-wrong-regular-node
			inputs: `={{
			((parameters) => {
				const mode = parameters?.mode;
				const inputs = [{ displayName: "Embedding", type: "${NodeConnectionType.AiEmbedding}", required: true, maxConnections: 1}]

				if (['insert', 'load'].includes(mode)) {
					inputs.push({ displayName: "", type: "${NodeConnectionType.Main}"})
				}

				if (mode === 'retrieve-as-tool') {
					return inputs;
				}

				if (mode === 'insert') {
					inputs.push({ displayName: "Document", type: "${NodeConnectionType.AiDocument}", required: true, maxConnections: 1})
				}
				return inputs
			})($parameter)
		}}`,
			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}"}]
				}
				return [{ displayName: "", type: "${NodeConnectionType.Main}"}]
			})($parameter)
		}}`,
			properties: [
				{
					displayName: 'Operation Mode',
					name: 'mode',
					type: 'options',
					noDataExpression: true,
					default: 'retrieve',
					options: [
						{
							name: 'Get Many',
							value: 'load',
							description: 'Get many ranked documents from vector store for query',
							action: 'Get many ranked documents from vector store for query',
						},
						{
							name: 'Insert Documents',
							value: 'insert',
							description: 'Insert documents into vector store',
							action: 'Insert documents into vector store',
						},
						{
							name: 'Retrieve Documents (For Agent/Chain)',
							value: 'retrieve',
							description: 'Retrieve documents from vector store to be used with AI nodes',
							action: 'Retrieve documents from vector store to be used with AI nodes',
							outputConnectionType: NodeConnectionType.AiVectorStore,
						},
						{
							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 Agent as Tool',
							outputConnectionType: NodeConnectionType.AiTool,
						},
					],
				},
				{
					...getConnectionHintNoticeField([NodeConnectionType.AiRetriever]),
					displayOptions: {
						show: {
							mode: ['retrieve'],
						},
					},
				},
				{
					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
				{
					displayName: 'Prompt',
					name: 'prompt',
					type: 'string',
					default: '',
					required: true,
					description:
						'Search prompt to retrieve matching documents from the vector store using similarity-based ranking',
					displayOptions: {
						show: {
							mode: ['load'],
						},
					},
				},
				{
					displayName: 'Limit',
					name: 'topK',
					type: 'number',
					default: 4,
					description: 'Number of top results to fetch from vector store',
					displayOptions: {
						show: {
							mode: ['load', 'retrieve-as-tool'],
						},
					},
				},
				{
					displayName: 'Include Metadata',
					name: 'includeDocumentMetadata',
					type: 'boolean',
					default: true,
					description: 'Whether or not to include document metadata',
					displayOptions: {
						show: {
							mode: ['retrieve-as-tool'],
						},
					},
				},
				...transformDescriptionForOperationMode(args.loadFields ?? [], 'load'),
				...transformDescriptionForOperationMode(args.retrieveFields ?? [], 'retrieve-as-tool'),
				...transformDescriptionForOperationMode(args.retrieveFields ?? [], 'retrieve'),
			],
		};

		methods = args.methods;

		async execute(this: IExecuteFunctions): Promise<INodeExecutionData[][]> {
			const mode = this.getNodeParameter('mode', 0) as 'load' | 'insert' | 'retrieve' | 'retrieve-as-tool';

			const embeddings = (await this.getInputConnectionData(
				NodeConnectionType.AiEmbedding,
				0,
			)) as Embeddings;

			if (mode === 'load') {
				const items = this.getInputData(0);

				const resultData = [];
				for (let itemIndex = 0; itemIndex < items.length; itemIndex++) {
					const filter = getMetadataFiltersValues(this, itemIndex);
					const vectorStore = await args.getVectorStoreClient(
						this,
						// We'll pass filter to similaritySearchVectorWithScore instaed of getVectorStoreClient
						undefined,
						embeddings,
						itemIndex,
					);
					const prompt = this.getNodeParameter('prompt', itemIndex) as string;
					const topK = this.getNodeParameter('topK', itemIndex, 4) as number;

					const embeddedPrompt = await embeddings.embedQuery(prompt);
					const docs = await vectorStore.similaritySearchVectorWithScore(
						embeddedPrompt,
						topK,
						filter,
					);

					const serializedDocs = docs.map(([doc, score]) => {
						const document = {
							metadata: doc.metadata,
							pageContent: doc.pageContent,
						};

						return {
							json: { document, score },
							pairedItem: {
								item: itemIndex,
							},
						};
					});

					resultData.push(...serializedDocs);
					logAiEvent(this, 'ai-vector-store-searched', { query: prompt });
				}

				return [resultData];
			}

			if (mode === 'insert') {
				const items = this.getInputData();

				const documentInput = (await this.getInputConnectionData(
					NodeConnectionType.AiDocument,
					0,
				)) as N8nJsonLoader | N8nBinaryLoader | Array<Document<Record<string, unknown>>>;

				const resultData = [];
				for (let itemIndex = 0; itemIndex < items.length; itemIndex++) {
					const itemData = items[itemIndex];
					const { processedDocuments, serializedDocuments } = await processDocument(
						documentInput,
						itemData,
						itemIndex,
					);
					resultData.push(...serializedDocuments);

					try {
						await args.populateVectorStore(this, embeddings, processedDocuments, itemIndex);

						logAiEvent(this, 'ai-vector-store-populated');
					} catch (error) {
						throw error;
					}
				}

				return [resultData];
			}

			throw new NodeOperationError(
				this.getNode(),
				'Only the "load" and "insert" operation modes are supported with execute',
			);
		}

		async supplyData(this: ISupplyDataFunctions, itemIndex: number): Promise<SupplyData> {
			const mode = this.getNodeParameter('mode', 0) as 'load' | 'insert' | 'retrieve' | 'retrieve-as-tool';
			const filter = getMetadataFiltersValues(this, itemIndex);
			const embeddings = (await this.getInputConnectionData(
				NodeConnectionType.AiEmbedding,
				0,
			)) as Embeddings;

			if (mode === 'retrieve') {
				const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex);
				return {
					response: logWrapper(vectorStore, this),
				};
			}

			if (mode === 'retrieve-as-tool') {
				const toolDescription = this.getNodeParameter('toolDescription', itemIndex) as string;
				const topK = this.getNodeParameter('topK', itemIndex, 4) as number;
				const node = this.getNode();
				const toolName = nodeNameToToolName(node);

				const searchClient = await args.getSearchClient(
					this,
					itemIndex
				)

				const vectorStoreTool = new DynamicTool({
					name: toolName,
					description: toolDescription,
					func: async (input) => {
						try {
							const embeddedPrompt = await embeddings.embedQuery(input);

							const searchResults = (await searchClient.search(
								input,
								{
									vectorSearchOptions: {
										queries: [
											{
												kind: 'vector',
												fields: ['contentVector', 'additionalContentVector'],
												kNearestNeighborsCount: 100,
												vector: embeddedPrompt,
											},
										],
										filterMode: 'preFilter'
									},
									top: topK,
									select: ['source', 'content', 'category', 'sourcePage', 'url', 'classification', 'id', 'title', 'documentNumber'],
									semanticSearchOptions: {
										configurationName: 'default',
										errorMode: 'partial'
									},
									queryType: 'semantic',
								}
							))as unknown as PagedAsyncIterableIterator<SearchResult<SearchDocument>>;


							const documents = [];

							for await (const result of searchResults) {
								documents.push(
									{
										type: 'text',
										text: JSON.stringify( result.document.content ?? "" ),
									}
								);
							}
							return documents;
						} catch (error) {
							throw new NodeOperationError(
								this.getNode(),
								`Error searching vector store: ${error.message}`,
							);
						}
					},
				});

				return {
					response: logWrapper(vectorStoreTool, this),
				};
			}

			throw new NodeOperationError(
				this.getNode(),
				'Only the "retrieve" operation mode is supported to supply data',
			);
		}
	};

VectorStoreAzureAISearch.node.ts:

import { AzureAISearchVectorStore, AzureAISearchConfig } from "@langchain/community/vectorstores/azure_aisearch";
import { SearchClient, AzureKeyCredential } from "@azure/search-documents";
import type { INodeProperties } from 'n8n-workflow';
import { createVectorStoreNode } from '../shared/createVectorStoreNode';
import { metadataFilterField } from '../../../utils/sharedFields';

const sharedFields: INodeProperties[] = [
	{
		displayName: 'Index Name',
		name: 'indexName',
		type: 'string',
		default: 'vectors',
		description: 'The Azure AI Search index name to store the vectors in',
	},
];

const insertFields: INodeProperties[] = [
	{
		displayName: 'Documents',
		name: 'documents',
		type: 'fixedCollection',
		default: {},
		placeholder: 'Add Document',
		typeOptions: {
			multipleValues: true,
		},
		options: [
			{
				name: 'documentValues',
				displayName: 'Document',
				values: [
					{
						displayName: 'Text',
						name: 'text',
						type: 'string',
						default: '',
						description: 'Text to be vectorized and stored',
					},
					{
						displayName: 'Metadata',
						name: 'metadata',
						type: 'json',
						default: '{}',
						description: 'Metadata to store with the document',
					},
				],
			},
		],
	},
];

const retrieveFields: INodeProperties[] = [
	{
		displayName: 'Query',
		name: 'query',
		type: 'string',
		default: '',
		description: 'The text to find similar documents to',
	},
	metadataFilterField,
];

// ... existing field definitions can be reused ...

export const VectorStoreAzureAISearch = createVectorStoreNode({
	meta: {
		description: 'Work with your data in Azure AI Search for vector-based search',
		icon: 'file:azure.svg',
		displayName: 'Azure AI Search Vector Store',
		name: 'vectorStoreAzureAISearch',
		credentials: [{ name: 'azureAISearchApi', required: true }],
		docsUrl: 'https://docs.n8n.io/...',
	},
	sharedFields,
	insertFields,
	loadFields: retrieveFields,
	retrieveFields,
	async getVectorStoreClient(context, filter, embeddings, itemIndex) {
		const indexName = context.getNodeParameter('indexName', itemIndex, '', {
			extractValue: true,
		}) as string;

		const credentials = await context.getCredentials('azureAISearchApi');

		const config: AzureAISearchConfig = {
			endpoint: String(credentials.endpoint),
			key: String(credentials.apiKey),
			indexName,
		};

		return new AzureAISearchVectorStore(embeddings, config);
	},
	async getSearchClient(context, itemIndex) {
		const indexName = context.getNodeParameter('indexName', itemIndex, '', {
			extractValue: true,
		}) as string;

		const credentials = await context.getCredentials('azureAISearchApi');

		return new SearchClient(
			String(credentials.endpoint),
			indexName,
			new AzureKeyCredential(String(credentials.apiKey))
		);
	},
	async populateVectorStore(context, embeddings, documents, itemIndex) {
		const indexName = context.getNodeParameter('indexName', itemIndex, '', {
			extractValue: true,
		}) as string;

		const credentials = await context.getCredentials('azureAISearchApi');

		const config = {
			endpoint: String(credentials.endpoint),
			key: String(credentials.apiKey),
			indexName,
		};

		await AzureAISearchVectorStore.fromDocuments(documents, embeddings, config);
	},
});

Other than this, helper.ts and logWrapper.ts were updated to accomodate logAiEvent and ISupplyDataFunctions.

Workflow: