fix(Postgres PGVector Store Node): Release postgres connections back to the pool (#12723)

Co-authored-by: Oleg Ivaniv <me@olegivaniv.com>
This commit is contained in:
कारतोफ्फेलस्क्रिप्ट™ 2025-01-27 11:52:18 +01:00 committed by GitHub
parent 02df25c450
commit 663dfb48de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 102 additions and 72 deletions

View file

@ -1,3 +1,4 @@
import type { MemoryVectorStore } from 'langchain/vectorstores/memory';
import type { INodeProperties } from 'n8n-workflow'; import type { INodeProperties } from 'n8n-workflow';
import { createVectorStoreNode } from '../shared/createVectorStoreNode'; import { createVectorStoreNode } from '../shared/createVectorStoreNode';
@ -20,7 +21,7 @@ const insertFields: INodeProperties[] = [
}, },
]; ];
export class VectorStoreInMemory extends createVectorStoreNode({ export class VectorStoreInMemory extends createVectorStoreNode<MemoryVectorStore>({
meta: { meta: {
displayName: 'In-Memory Vector Store', displayName: 'In-Memory Vector Store',
name: 'vectorStoreInMemory', name: 'vectorStoreInMemory',

View file

@ -213,7 +213,7 @@ class ExtendedPGVectorStore extends PGVectorStore {
} }
} }
export class VectorStorePGVector extends createVectorStoreNode({ export class VectorStorePGVector extends createVectorStoreNode<ExtendedPGVectorStore>({
meta: { meta: {
description: 'Work with your data in Postgresql with the PGVector extension', description: 'Work with your data in Postgresql with the PGVector extension',
icon: 'file:postgres.svg', icon: 'file:postgres.svg',
@ -274,6 +274,7 @@ export class VectorStorePGVector extends createVectorStoreNode({
return await ExtendedPGVectorStore.initialize(embeddings, config); return await ExtendedPGVectorStore.initialize(embeddings, config);
}, },
async populateVectorStore(context, embeddings, documents, itemIndex) { async populateVectorStore(context, embeddings, documents, itemIndex) {
// NOTE: if you are to create the HNSW index before use, you need to consider moving the distanceStrategy field to // NOTE: if you are to create the HNSW index before use, you need to consider moving the distanceStrategy field to
// shared fields, because you need that strategy when creating the index. // shared fields, because you need that strategy when creating the index.
@ -307,6 +308,11 @@ export class VectorStorePGVector extends createVectorStoreNode({
metadataColumnName: 'metadata', metadataColumnName: 'metadata',
}) as ColumnOptions; }) as ColumnOptions;
await PGVectorStore.fromDocuments(documents, embeddings, config); const vectorStore = await PGVectorStore.fromDocuments(documents, embeddings, config);
vectorStore.client?.release();
},
releaseVectorStoreClient(vectorStore) {
vectorStore.client?.release();
}, },
}) {} }) {}

View file

@ -51,7 +51,7 @@ const insertFields: INodeProperties[] = [
}, },
]; ];
export class VectorStorePinecone extends createVectorStoreNode({ export class VectorStorePinecone extends createVectorStoreNode<PineconeStore>({
meta: { meta: {
displayName: 'Pinecone Vector Store', displayName: 'Pinecone Vector Store',
name: 'vectorStorePinecone', name: 'vectorStorePinecone',

View file

@ -79,7 +79,7 @@ const retrieveFields: INodeProperties[] = [
}, },
]; ];
export class VectorStoreQdrant extends createVectorStoreNode({ export class VectorStoreQdrant extends createVectorStoreNode<ExtendedQdrantVectorStore>({
meta: { meta: {
displayName: 'Qdrant Vector Store', displayName: 'Qdrant Vector Store',
name: 'vectorStoreQdrant', name: 'vectorStoreQdrant',

View file

@ -41,7 +41,7 @@ const retrieveFields: INodeProperties[] = [
const updateFields: INodeProperties[] = [...insertFields]; const updateFields: INodeProperties[] = [...insertFields];
export class VectorStoreSupabase extends createVectorStoreNode({ export class VectorStoreSupabase extends createVectorStoreNode<SupabaseVectorStore>({
meta: { meta: {
description: 'Work with your data in Supabase Vector Store', description: 'Work with your data in Supabase Vector Store',
icon: 'file:supabase.svg', icon: 'file:supabase.svg',

View file

@ -46,7 +46,7 @@ const retrieveFields: INodeProperties[] = [
}, },
]; ];
export class VectorStoreZep extends createVectorStoreNode({ export class VectorStoreZep extends createVectorStoreNode<ZepVectorStore | ZepCloudVectorStore>({
meta: { meta: {
displayName: 'Zep Vector Store', displayName: 'Zep Vector Store',
name: 'vectorStoreZep', name: 'vectorStoreZep',

View file

@ -49,7 +49,7 @@ interface NodeMeta {
operationModes?: NodeOperationMode[]; operationModes?: NodeOperationMode[];
} }
export interface VectorStoreNodeConstructorArgs { export interface VectorStoreNodeConstructorArgs<T extends VectorStore = VectorStore> {
meta: NodeMeta; meta: NodeMeta;
methods?: { methods?: {
listSearch?: { listSearch?: {
@ -77,7 +77,8 @@ export interface VectorStoreNodeConstructorArgs {
filter: Record<string, never> | undefined, filter: Record<string, never> | undefined,
embeddings: Embeddings, embeddings: Embeddings,
itemIndex: number, itemIndex: number,
) => Promise<VectorStore>; ) => Promise<T>;
releaseVectorStoreClient?: (vectorStore: T) => void;
} }
function transformDescriptionForOperationMode( function transformDescriptionForOperationMode(
@ -90,11 +91,15 @@ function transformDescriptionForOperationMode(
})); }));
} }
function isUpdateSupported(args: VectorStoreNodeConstructorArgs): boolean { function isUpdateSupported<T extends VectorStore>(
args: VectorStoreNodeConstructorArgs<T>,
): boolean {
return args.meta.operationModes?.includes('update') ?? false; return args.meta.operationModes?.includes('update') ?? false;
} }
function getOperationModeOptions(args: VectorStoreNodeConstructorArgs): INodePropertyOptions[] { function getOperationModeOptions<T extends VectorStore>(
args: VectorStoreNodeConstructorArgs<T>,
): INodePropertyOptions[] {
const enabledOperationModes = args.meta.operationModes ?? DEFAULT_OPERATION_MODES; const enabledOperationModes = args.meta.operationModes ?? DEFAULT_OPERATION_MODES;
const allOptions = [ const allOptions = [
@ -137,7 +142,9 @@ function getOperationModeOptions(args: VectorStoreNodeConstructorArgs): INodePro
); );
} }
export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) => export const createVectorStoreNode = <T extends VectorStore = VectorStore>(
args: VectorStoreNodeConstructorArgs<T>,
) =>
class VectorStoreNodeType implements INodeType { class VectorStoreNodeType implements INodeType {
description: INodeTypeDescription = { description: INodeTypeDescription = {
displayName: args.meta.displayName, displayName: args.meta.displayName,
@ -334,6 +341,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
embeddings, embeddings,
itemIndex, itemIndex,
); );
try {
const prompt = this.getNodeParameter('prompt', itemIndex) as string; const prompt = this.getNodeParameter('prompt', itemIndex) as string;
const topK = this.getNodeParameter('topK', itemIndex, 4) as number; const topK = this.getNodeParameter('topK', itemIndex, 4) as number;
@ -366,6 +374,9 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
resultData.push(...serializedDocs); resultData.push(...serializedDocs);
logAiEvent(this, 'ai-vector-store-searched', { query: prompt }); logAiEvent(this, 'ai-vector-store-searched', { query: prompt });
} finally {
args.releaseVectorStoreClient?.(vectorStore);
}
} }
return [resultData]; return [resultData];
@ -427,6 +438,7 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
itemIndex, itemIndex,
); );
try {
const { processedDocuments, serializedDocuments } = await processDocument( const { processedDocuments, serializedDocuments } = await processDocument(
loader, loader,
itemData, itemData,
@ -445,6 +457,9 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
}); });
logAiEvent(this, 'ai-vector-store-updated'); logAiEvent(this, 'ai-vector-store-updated');
} finally {
args.releaseVectorStoreClient?.(vectorStore);
}
} }
return [resultData]; return [resultData];
@ -468,6 +483,9 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex); const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex);
return { return {
response: logWrapper(vectorStore, this), response: logWrapper(vectorStore, this),
closeFunction: async () => {
args.releaseVectorStoreClient?.(vectorStore);
},
}; };
} }
@ -491,6 +509,8 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
embeddings, embeddings,
itemIndex, itemIndex,
); );
try {
const embeddedPrompt = await embeddings.embedQuery(input); const embeddedPrompt = await embeddings.embedQuery(input);
const documents = await vectorStore.similaritySearchVectorWithScore( const documents = await vectorStore.similaritySearchVectorWithScore(
embeddedPrompt, embeddedPrompt,
@ -508,6 +528,9 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
}; };
}) })
.filter((document) => !!document); .filter((document) => !!document);
} finally {
args.releaseVectorStoreClient?.(vectorStore);
}
}, },
}); });