diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.test.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.test.ts index 40fc1d494d1..43bfd892ee8 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.test.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.test.ts @@ -5,7 +5,12 @@ import type { Embeddings } from '@langchain/core/embeddings'; import type { VectorStore } from '@langchain/core/vectorstores'; import { mock } from 'jest-mock-extended'; import type { DynamicTool } from 'langchain/tools'; -import type { ISupplyDataFunctions, NodeParameterValueType } from 'n8n-workflow'; +import type { + IExecuteFunctions, + ISupplyDataFunctions, + NodeParameterValueType, + INodeExecutionData, +} from 'n8n-workflow'; import { createVectorStoreNode } from './createVectorStoreNode'; import type { VectorStoreNodeConstructorArgs } from './types'; @@ -16,6 +21,7 @@ jest.mock('@utils/logWrapper', () => ({ const DEFAULT_PARAMETERS = { options: {}, + useReranker: false, topK: 1, }; @@ -222,4 +228,177 @@ describe('createVectorStoreNode', () => { ]); }); }); + + describe('execute mode', () => { + const executeContext = mock({ + getNodeParameter: jest.fn(), + getInputConnectionData: jest.fn().mockReturnValue(embeddings), + getInputData: jest.fn(), + }); + + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('retrieve-as-tool mode in execute context', () => { + it('should execute retrieve-as-tool and return documents with metadata', async () => { + // ARRANGE + const parameters: Record = { + ...DEFAULT_PARAMETERS, + mode: 'retrieve-as-tool', + includeDocumentMetadata: true, + }; + + const inputData: INodeExecutionData[] = [ + { + json: { input: MOCK_SEARCH_VALUE }, + pairedItem: { item: 0 }, + }, + ]; + + executeContext.getNodeParameter.mockImplementation( + (parameterName: string): NodeParameterValueType | object => parameters[parameterName], + ); + executeContext.getInputData.mockReturnValue(inputData); + + // ACT + const VectorStoreNodeType = createVectorStoreNode(vectorStoreNodeArgs); + const nodeType = new VectorStoreNodeType(); + const result = await nodeType.execute.call(executeContext); + + // ASSERT + expect(result).toHaveLength(1); // One output array + expect(result[0][0]?.json?.response).toHaveLength(2); // Two documents returned + + expect(result[0][0]).toEqual({ + json: { + response: [ + { + type: 'text', + text: JSON.stringify({ + pageContent: 'first page', + metadata: { id: 123 }, + }), + }, + { + type: 'text', + text: JSON.stringify({ + pageContent: 'second page', + metadata: { id: 567 }, + }), + }, + ], + }, + pairedItem: { item: 0 }, + }); + + expect(embeddings.embedQuery).toHaveBeenCalledWith(MOCK_SEARCH_VALUE); + expect(vectorStore.similaritySearchVectorWithScore).toHaveBeenCalledWith( + MOCK_EMBEDDED_SEARCH_VALUE, + parameters.topK, + undefined, // filter + ); + }); + + it('should execute retrieve-as-tool and return documents without metadata', async () => { + // ARRANGE + const parameters: Record = { + ...DEFAULT_PARAMETERS, + mode: 'retrieve-as-tool', + includeDocumentMetadata: false, + }; + + const inputData: INodeExecutionData[] = [ + { + json: { input: MOCK_SEARCH_VALUE }, + pairedItem: { item: 0 }, + }, + ]; + + executeContext.getNodeParameter.mockImplementation( + (parameterName: string): NodeParameterValueType | object => parameters[parameterName], + ); + executeContext.getInputData.mockReturnValue(inputData); + + // ACT + const VectorStoreNodeType = createVectorStoreNode(vectorStoreNodeArgs); + const nodeType = new VectorStoreNodeType(); + const result = await nodeType.execute.call(executeContext); + + // ASSERT + expect(result[0][0].json.response).toHaveLength(2); + const response = result[0][0].json.response as Array<{ pageContent: string }>; + const doc0 = JSON.parse(response[0].pageContent); + const doc1 = JSON.parse(response[1].pageContent); + expect(doc0).not.toHaveProperty('metadata'); + expect(doc0).toEqual({ pageContent: 'first page' }); + expect(doc1).toEqual({ pageContent: 'second page' }); + }); + + it('should process multiple input items', async () => { + // ARRANGE + const parameters: Record = { + ...DEFAULT_PARAMETERS, + mode: 'retrieve-as-tool', + includeDocumentMetadata: true, + }; + + const inputData: INodeExecutionData[] = [ + { + json: { input: 'first query' }, + pairedItem: { item: 0 }, + }, + { + json: { input: 'second query' }, + pairedItem: { item: 1 }, + }, + ]; + + executeContext.getNodeParameter.mockImplementation( + (parameterName: string): NodeParameterValueType | object => parameters[parameterName], + ); + executeContext.getInputData.mockReturnValue(inputData); + + // ACT + const VectorStoreNodeType = createVectorStoreNode(vectorStoreNodeArgs); + const nodeType = new VectorStoreNodeType(); + const result = await nodeType.execute.call(executeContext); + + // ASSERT + expect(result).toHaveLength(1); + expect(result[0]).toHaveLength(2); // One result item per input query + + // Check that embedQuery was called for both input queries + expect(embeddings.embedQuery).toHaveBeenCalledTimes(2); + expect(embeddings.embedQuery).toHaveBeenNthCalledWith(1, 'first query'); + expect(embeddings.embedQuery).toHaveBeenNthCalledWith(2, 'second query'); + + // Check pairedItem references and that each result contains both documents + expect(result[0][0].pairedItem).toEqual({ item: 0 }); + expect(result[0][0].json.response).toHaveLength(2); // 2 documents for first query + expect(result[0][1].pairedItem).toEqual({ item: 1 }); + expect(result[0][1].json.response).toHaveLength(2); // 2 documents for second query + }); + + it('should throw error for unsupported mode in execute', async () => { + // ARRANGE + const parameters: Record = { + ...DEFAULT_PARAMETERS, + mode: 'retrieve', // This mode is not supported in execute + }; + + executeContext.getNodeParameter.mockImplementation( + (parameterName: string): NodeParameterValueType | object => parameters[parameterName], + ); + + // ACT & ASSERT + const VectorStoreNodeType = createVectorStoreNode(vectorStoreNodeArgs); + const nodeType = new VectorStoreNodeType(); + + await expect(nodeType.execute.call(executeContext)).rejects.toThrow( + 'Only the "load", "update", "insert", and "retrieve-as-tool" operation modes are supported with execute', + ); + }); + }); + }); }); diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.ts index 7fff59b5203..3cbd58bf91d 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/createVectorStoreNode.ts @@ -20,6 +20,7 @@ import { handleUpdateOperation, handleRetrieveOperation, handleRetrieveAsToolOperation, + handleRetrieveAsToolExecuteOperation, } from './operations'; import type { NodeOperationMode, VectorStoreNodeConstructorArgs } from './types'; // Import utility functions @@ -291,9 +292,26 @@ export const createVectorStoreNode = ( return [resultData]; } + if (mode === 'retrieve-as-tool') { + const items = this.getInputData(0); + const resultData = []; + + for (let itemIndex = 0; itemIndex < items.length; itemIndex++) { + const docs = await handleRetrieveAsToolExecuteOperation( + this, + args, + embeddings, + itemIndex, + ); + resultData.push(...docs); + } + + return [resultData]; + } + throw new NodeOperationError( this.getNode(), - 'Only the "load", "update" and "insert" operation modes are supported with execute', + 'Only the "load", "update", "insert", and "retrieve-as-tool" operation modes are supported with execute', ); } diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/__tests__/retrieveAsToolExecuteOperation.test.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/__tests__/retrieveAsToolExecuteOperation.test.ts new file mode 100644 index 00000000000..608a6868b87 --- /dev/null +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/__tests__/retrieveAsToolExecuteOperation.test.ts @@ -0,0 +1,425 @@ +/* eslint-disable @typescript-eslint/no-unsafe-return */ +/* eslint-disable @typescript-eslint/unbound-method */ +import type { Document } from '@langchain/core/documents'; +import type { Embeddings } from '@langchain/core/embeddings'; +import type { BaseDocumentCompressor } from '@langchain/core/retrievers/document_compressors'; +import type { VectorStore } from '@langchain/core/vectorstores'; +import type { MockProxy } from 'jest-mock-extended'; +import { mock } from 'jest-mock-extended'; +import type { IExecuteFunctions, INodeExecutionData } from 'n8n-workflow'; +import { NodeConnectionTypes } from 'n8n-workflow'; + +import { logAiEvent } from '@utils/helpers'; + +import type { VectorStoreNodeConstructorArgs } from '../../types'; +import { handleRetrieveAsToolExecuteOperation } from '../retrieveAsToolExecuteOperation'; + +// Mock helper functions from external modules +jest.mock('@utils/helpers', () => ({ + getMetadataFiltersValues: jest.fn().mockReturnValue({ testFilter: 'value' }), + logAiEvent: jest.fn(), +})); + +describe('handleRetrieveAsToolExecuteOperation', () => { + let mockContext: MockProxy; + let mockEmbeddings: MockProxy; + let mockVectorStore: MockProxy; + let mockReranker: MockProxy; + let mockArgs: VectorStoreNodeConstructorArgs; + let nodeParameters: Record; + let inputData: INodeExecutionData[]; + + beforeEach(() => { + nodeParameters = { + topK: 3, + includeDocumentMetadata: true, + useReranker: false, + }; + + inputData = [ + { + json: { input: 'test search query' }, + pairedItem: { item: 0 }, + }, + ]; + + mockContext = mock(); + mockContext.getNodeParameter.mockImplementation((parameterName, _itemIndex, fallbackValue) => { + if (typeof parameterName !== 'string') return fallbackValue; + return nodeParameters[parameterName] ?? fallbackValue; + }); + mockContext.getInputData.mockReturnValue(inputData); + + mockEmbeddings = mock(); + mockEmbeddings.embedQuery.mockResolvedValue([0.1, 0.2, 0.3]); + + mockVectorStore = mock(); + mockVectorStore.similaritySearchVectorWithScore.mockResolvedValue([ + [{ pageContent: 'test content 1', metadata: { test: 'metadata 1' } } as Document, 0.95], + [{ pageContent: 'test content 2', metadata: { test: 'metadata 2' } } as Document, 0.85], + [{ pageContent: 'test content 3', metadata: { test: 'metadata 3' } } as Document, 0.75], + ]); + + mockReranker = mock(); + mockReranker.compressDocuments.mockResolvedValue([ + { + pageContent: 'test content 2', + metadata: { test: 'metadata 2', relevanceScore: 0.98 }, + } as Document, + { + pageContent: 'test content 1', + metadata: { test: 'metadata 1', relevanceScore: 0.92 }, + } as Document, + { + pageContent: 'test content 3', + metadata: { test: 'metadata 3', relevanceScore: 0.88 }, + } as Document, + ]); + + mockContext.getInputConnectionData.mockResolvedValue(mockReranker); + + mockArgs = { + meta: { + displayName: 'Test Vector Store', + name: 'testVectorStore', + description: 'Vector store for testing', + docsUrl: 'https://example.com', + icon: 'file:testIcon.svg', + }, + sharedFields: [], + getVectorStoreClient: jest.fn().mockResolvedValue(mockVectorStore), + populateVectorStore: jest.fn().mockResolvedValue(undefined), + releaseVectorStoreClient: jest.fn(), + }; + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should retrieve documents from vector store using query from input data', async () => { + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(mockArgs.getVectorStoreClient).toHaveBeenCalledWith( + mockContext, + undefined, + mockEmbeddings, + 0, + ); + + expect(mockEmbeddings.embedQuery).toHaveBeenCalledWith('test search query'); + + expect(mockVectorStore.similaritySearchVectorWithScore).toHaveBeenCalledWith( + [0.1, 0.2, 0.3], + 3, + { testFilter: 'value' }, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + const response = result[0].json.response as Array<{ type: string; text: string }>; + expect(response[0]).toEqual({ + type: 'text', + text: JSON.stringify({ + pageContent: 'test content 1', + metadata: { test: 'metadata 1' }, + }), + }); + expect(result[0].pairedItem).toEqual({ item: 0 }); + + expect(mockArgs.releaseVectorStoreClient).toHaveBeenCalledWith(mockVectorStore); + expect(logAiEvent).toHaveBeenCalledWith(mockContext, 'ai-vector-store-searched', { + input: 'test search query', + }); + }); + + it('should throw error when input data does not contain query', async () => { + inputData[0].json = { notQuery: 'some value' }; + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Input data must contain a "input" field with the search query'); + }); + + it('should throw error when query is not a string', async () => { + inputData[0].json = { input: 123 }; + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Input data must contain a "input" field with the search query'); + }); + + it('should throw error when query is empty string', async () => { + inputData[0].json = { input: '' }; + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Input data must contain a "input" field with the search query'); + }); + + it('should include metadata when includeDocumentMetadata is true', async () => { + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + const response = result[0].json.response as Array<{ type: string; text: string }>; + const firstDoc = JSON.parse(response[0].text); + expect(firstDoc).toHaveProperty('metadata'); + expect(firstDoc.metadata).toEqual({ test: 'metadata 1' }); + }); + + it('should exclude metadata when includeDocumentMetadata is false', async () => { + nodeParameters.includeDocumentMetadata = false; + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + const response = result[0].json.response as Array<{ pageContent: string }>; + const firstDoc = JSON.parse(response[0].pageContent); + expect(firstDoc).not.toHaveProperty('metadata'); + expect(firstDoc).toEqual({ + pageContent: 'test content 1', + }); + }); + + it('should limit results based on topK parameter', async () => { + nodeParameters.topK = 1; + + await handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0); + + expect(mockVectorStore.similaritySearchVectorWithScore).toHaveBeenCalledWith( + expect.anything(), + 1, + expect.anything(), + ); + }); + + it('should use topK default value when not provided', async () => { + delete nodeParameters.topK; + + await handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0); + + expect(mockVectorStore.similaritySearchVectorWithScore).toHaveBeenCalledWith( + expect.anything(), + 4, // default value + expect.anything(), + ); + }); + + it('should release vector store client even if search fails', async () => { + mockVectorStore.similaritySearchVectorWithScore.mockRejectedValueOnce( + new Error('Search failed'), + ); + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Search failed'); + + expect(mockArgs.releaseVectorStoreClient).toHaveBeenCalledWith(mockVectorStore); + }); + + describe('reranking functionality', () => { + beforeEach(() => { + nodeParameters.useReranker = true; + }); + + it('should use reranker when useReranker is true', async () => { + await handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0); + + expect(mockContext.getInputConnectionData).toHaveBeenCalledWith( + NodeConnectionTypes.AiReranker, + 0, + ); + expect(mockReranker.compressDocuments).toHaveBeenCalledWith( + [ + { pageContent: 'test content 1', metadata: { test: 'metadata 1' } }, + { pageContent: 'test content 2', metadata: { test: 'metadata 2' } }, + { pageContent: 'test content 3', metadata: { test: 'metadata 3' } }, + ], + 'test search query', + ); + }); + + it('should return reranked documents in the correct order', async () => { + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + const response = result[0].json.response as Array<{ type: string; text: string }>; + + // First result should be the reranked first document (was second in original order) + const doc0 = JSON.parse(response[0].text); + expect(doc0.pageContent).toEqual('test content 2'); + expect(doc0.metadata).toEqual({ test: 'metadata 2' }); + + // Second result should be the reranked second document (was first in original order) + const doc1 = JSON.parse(response[1].text); + expect(doc1.pageContent).toEqual('test content 1'); + expect(doc1.metadata).toEqual({ test: 'metadata 1' }); + + // Third result should be the reranked third document + const doc2 = JSON.parse(response[2].text); + expect(doc2.pageContent).toEqual('test content 3'); + expect(doc2.metadata).toEqual({ test: 'metadata 3' }); + }); + + it('should handle reranking with includeDocumentMetadata false', async () => { + nodeParameters.includeDocumentMetadata = false; + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + const response = result[0].json.response as Array<{ pageContent: string }>; + + // Should maintain reranked order but exclude metadata + const doc0 = JSON.parse(response[0].pageContent); + expect(doc0).toEqual({ pageContent: 'test content 2' }); + const doc1 = JSON.parse(response[1].pageContent); + expect(doc1).toEqual({ pageContent: 'test content 1' }); + const doc2 = JSON.parse(response[2].pageContent); + expect(doc2).toEqual({ pageContent: 'test content 3' }); + }); + + it('should not call reranker when useReranker is false', async () => { + nodeParameters.useReranker = false; + + await handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0); + + expect(mockContext.getInputConnectionData).not.toHaveBeenCalled(); + expect(mockReranker.compressDocuments).not.toHaveBeenCalled(); + }); + + it('should release vector store client even if reranking fails', async () => { + mockReranker.compressDocuments.mockRejectedValueOnce(new Error('Reranking failed')); + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Reranking failed'); + + expect(mockArgs.releaseVectorStoreClient).toHaveBeenCalledWith(mockVectorStore); + }); + + it('should properly handle relevanceScore from reranker metadata', async () => { + // Mock reranker to return documents with relevanceScore in different metadata structure + mockReranker.compressDocuments.mockResolvedValueOnce([ + { + pageContent: 'test content 2', + metadata: { test: 'metadata 2', relevanceScore: 0.98, otherField: 'value' }, + } as Document, + { + pageContent: 'test content 1', + metadata: { test: 'metadata 1', relevanceScore: 0.92 }, + } as Document, + ]); + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(2); + const response = result[0].json.response as Array<{ type: string; text: string }>; + + // Check that relevanceScore is used but metadata is preserved without relevanceScore + const doc0 = JSON.parse(response[0].text); + expect(doc0.metadata).toEqual({ test: 'metadata 2', otherField: 'value' }); + expect(doc0.metadata).not.toHaveProperty('relevanceScore'); + + const doc1 = JSON.parse(response[1].text); + expect(doc1.metadata).toEqual({ test: 'metadata 1' }); + expect(doc1.metadata).not.toHaveProperty('relevanceScore'); + }); + + it('should not use reranker when no documents are found', async () => { + mockVectorStore.similaritySearchVectorWithScore.mockResolvedValueOnce([]); + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(mockContext.getInputConnectionData).not.toHaveBeenCalled(); + expect(mockReranker.compressDocuments).not.toHaveBeenCalled(); + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(0); + }); + }); + + describe('empty result handling', () => { + it('should return empty array when vector store returns no documents', async () => { + mockVectorStore.similaritySearchVectorWithScore.mockResolvedValueOnce([]); + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(0); + expect(logAiEvent).toHaveBeenCalledWith(mockContext, 'ai-vector-store-searched', { + input: 'test search query', + }); + }); + }); + + describe('error handling', () => { + it('should release client resources when embedQuery fails', async () => { + mockEmbeddings.embedQuery.mockRejectedValueOnce(new Error('Embedding failed')); + + await expect( + handleRetrieveAsToolExecuteOperation(mockContext, mockArgs, mockEmbeddings, 0), + ).rejects.toThrow('Embedding failed'); + + expect(mockArgs.releaseVectorStoreClient).toHaveBeenCalledWith(mockVectorStore); + }); + + it('should handle missing releaseVectorStoreClient function gracefully', async () => { + delete mockArgs.releaseVectorStoreClient; + + const result = await handleRetrieveAsToolExecuteOperation( + mockContext, + mockArgs, + mockEmbeddings, + 0, + ); + + expect(result).toHaveLength(1); + expect(result[0].json.response).toHaveLength(3); + // Should not throw error when releaseVectorStoreClient is undefined + }); + }); +}); diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/index.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/index.ts index 74d2c6cd41a..e42f1897733 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/index.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/index.ts @@ -3,3 +3,4 @@ export * from './insertOperation'; export * from './updateOperation'; export * from './retrieveOperation'; export * from './retrieveAsToolOperation'; +export * from './retrieveAsToolExecuteOperation'; diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/retrieveAsToolExecuteOperation.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/retrieveAsToolExecuteOperation.ts new file mode 100644 index 00000000000..84e37bc70a0 --- /dev/null +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/shared/createVectorStoreNode/operations/retrieveAsToolExecuteOperation.ts @@ -0,0 +1,111 @@ +import type { Embeddings } from '@langchain/core/embeddings'; +import type { BaseDocumentCompressor } from '@langchain/core/retrievers/document_compressors'; +import type { VectorStore } from '@langchain/core/vectorstores'; +import { + assertParamIsBoolean, + assertParamIsNumber, + NodeConnectionTypes, + type IExecuteFunctions, + type INodeExecutionData, +} from 'n8n-workflow'; + +import { getMetadataFiltersValues, logAiEvent } from '@utils/helpers'; + +import type { VectorStoreNodeConstructorArgs } from '../types'; + +/** + * Handles the 'retrieve-as-tool' operation mode in execute context + * Searches the vector store for documents similar to a query and returns execution data + * This is similar to the load operation but designed to work with the new tool execution pattern + */ +export async function handleRetrieveAsToolExecuteOperation( + context: IExecuteFunctions, + args: VectorStoreNodeConstructorArgs, + embeddings: Embeddings, + itemIndex: number, +): Promise { + const filter = getMetadataFiltersValues(context, itemIndex); + const vectorStore = await args.getVectorStoreClient( + context, + // We'll pass filter to similaritySearchVectorWithScore instead of getVectorStoreClient + undefined, + embeddings, + itemIndex, + ); + + try { + // Get the search parameters - query from input data, others from node parameters + const inputData = context.getInputData(); + const item = inputData[itemIndex]; + const query = typeof item.json.input === 'string' ? item.json.input : undefined; + + if (!query || typeof query !== 'string') { + throw new Error('Input data must contain a "input" field with the search query'); + } + + const topK = context.getNodeParameter('topK', itemIndex, 4); + assertParamIsNumber('topK', topK, context.getNode()); + const useReranker = context.getNodeParameter('useReranker', itemIndex, false); + assertParamIsBoolean('useReranker', useReranker, context.getNode()); + + const includeDocumentMetadata = context.getNodeParameter( + 'includeDocumentMetadata', + itemIndex, + true, + ); + assertParamIsBoolean('includeDocumentMetadata', includeDocumentMetadata, context.getNode()); + + // Embed the query to prepare for vector similarity search + const embeddedQuery = await embeddings.embedQuery(query); + + // Get the most similar documents to the embedded query + let docs = await vectorStore.similaritySearchVectorWithScore(embeddedQuery, topK, filter); + + // If reranker is used, rerank the documents + if (useReranker && docs.length > 0) { + const reranker = (await context.getInputConnectionData( + NodeConnectionTypes.AiReranker, + 0, + )) as BaseDocumentCompressor; + const documents = docs.map(([doc]) => doc); + + const rerankedDocuments = await reranker.compressDocuments(documents, query); + docs = rerankedDocuments.map((doc) => { + const { relevanceScore, ...metadata } = doc.metadata || {}; + return [{ ...doc, metadata }, relevanceScore ?? 0]; + }); + } + + // Format the documents for the output similar to the original tool format + const serializedDocs = docs.map(([doc]) => { + if (includeDocumentMetadata) { + return { + type: 'text', + text: JSON.stringify({ ...doc }), + }; + } else { + return { + type: 'text', + pageContent: JSON.stringify({ pageContent: doc.pageContent }), + }; + } + }); + + // Log the AI event for analytics + logAiEvent(context, 'ai-vector-store-searched', { input: query }); + + return [ + { + json: { + response: serializedDocs, + }, + pairedItem: { + item: itemIndex, + }, + }, + ]; + } finally { + // Release the vector store client if a release method was provided + args.releaseVectorStoreClient?.(vectorStore); + } +}