fix(Zep Vector Store Node): Cloud vector store integration (#12353)

This commit is contained in:
Eugene 2025-01-06 14:46:57 +01:00 committed by GitHub
parent 536a975534
commit 2433d6b7d3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 90 additions and 6 deletions

View file

@ -0,0 +1,73 @@
import { ZepVectorStore } from '@langchain/community/vectorstores/zep';
import { ZepCloudVectorStore } from '@langchain/community/vectorstores/zep_cloud';
import { mock } from 'jest-mock-extended';
import type { ISupplyDataFunctions } from 'n8n-workflow';
import { VectorStoreZep } from './VectorStoreZep.node';
describe('VectorStoreZep', () => {
const vectorStore = new VectorStoreZep();
const helpers = mock<ISupplyDataFunctions['helpers']>();
const executeFunctions = mock<ISupplyDataFunctions>({ helpers });
beforeEach(() => {
jest.resetAllMocks();
executeFunctions.addInputData.mockReturnValue({ index: 0 });
});
it('should get vector store cloud client', async () => {
executeFunctions.getNodeParameter.mockImplementation((paramName: string) => {
switch (paramName) {
case 'mode':
return 'retrieve';
case 'collectionName':
return 'test-collection';
case 'options':
return {};
default:
return undefined;
}
});
executeFunctions.getCredentials.mockResolvedValue(
mock({
apiKey: 'some-key',
cloud: true,
}),
);
const { response } = await vectorStore.supplyData.call(executeFunctions, 0);
expect(response).toBeDefined();
expect(response).toBeInstanceOf(ZepCloudVectorStore);
});
it('should get vector store self-hosted client', async () => {
executeFunctions.getNodeParameter.mockImplementation((paramName: string) => {
switch (paramName) {
case 'mode':
return 'retrieve';
case 'collectionName':
return 'test-collection';
case 'options':
return {};
default:
return undefined;
}
});
executeFunctions.getCredentials.mockResolvedValue(
mock({
apiKey: 'some-key',
apiUrl: 'https://example.com',
cloud: false,
}),
);
const { response } = await vectorStore.supplyData.call(executeFunctions, 0);
expect(response).toBeDefined();
expect(response).toBeInstanceOf(ZepVectorStore);
});
});

View file

@ -1,5 +1,5 @@
import type { IZepConfig } from '@langchain/community/vectorstores/zep';
import { ZepVectorStore } from '@langchain/community/vectorstores/zep';
import { ZepCloudVectorStore } from '@langchain/community/vectorstores/zep_cloud';
import type { IDataObject, INodeProperties } from 'n8n-workflow';
import { NodeOperationError } from 'n8n-workflow';
@ -84,17 +84,21 @@ export class VectorStoreZep extends createVectorStoreNode({
const credentials = await context.getCredentials<{
apiKey?: string;
apiUrl: string;
cloud: boolean;
}>('zepApi');
const zepConfig: IZepConfig = {
apiUrl: credentials.apiUrl,
const zepConfig = {
apiKey: credentials.apiKey,
collectionName,
embeddingDimensions: options.embeddingDimensions ?? 1536,
metadata: filter,
};
return new ZepVectorStore(embeddings, zepConfig);
if (credentials.cloud) {
return new ZepCloudVectorStore(embeddings, zepConfig);
} else {
return new ZepVectorStore(embeddings, { ...zepConfig, apiUrl: credentials.apiUrl });
}
},
async populateVectorStore(context, embeddings, documents, itemIndex) {
const collectionName = context.getNodeParameter('collectionName', itemIndex) as string;
@ -107,10 +111,10 @@ export class VectorStoreZep extends createVectorStoreNode({
const credentials = await context.getCredentials<{
apiKey?: string;
apiUrl: string;
cloud: boolean;
}>('zepApi');
const zepConfig = {
apiUrl: credentials.apiUrl,
apiKey: credentials.apiKey,
collectionName,
embeddingDimensions: options.embeddingDimensions ?? 1536,
@ -118,7 +122,14 @@ export class VectorStoreZep extends createVectorStoreNode({
};
try {
await ZepVectorStore.fromDocuments(documents, embeddings, zepConfig);
if (credentials.cloud) {
await ZepCloudVectorStore.fromDocuments(documents, embeddings, zepConfig);
} else {
await ZepVectorStore.fromDocuments(documents, embeddings, {
...zepConfig,
apiUrl: credentials.apiUrl,
});
}
} catch (error) {
const errorCode = (error as IDataObject).code as number;
const responseData = (error as IDataObject).responseData as string;