mirror of
https://github.com/n8n-io/n8n.git
synced 2025-03-05 20:50:17 -08:00
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:
parent
02df25c450
commit
663dfb48de
|
@ -1,3 +1,4 @@
|
|||
import type { MemoryVectorStore } from 'langchain/vectorstores/memory';
|
||||
import type { INodeProperties } from 'n8n-workflow';
|
||||
|
||||
import { createVectorStoreNode } from '../shared/createVectorStoreNode';
|
||||
|
@ -20,7 +21,7 @@ const insertFields: INodeProperties[] = [
|
|||
},
|
||||
];
|
||||
|
||||
export class VectorStoreInMemory extends createVectorStoreNode({
|
||||
export class VectorStoreInMemory extends createVectorStoreNode<MemoryVectorStore>({
|
||||
meta: {
|
||||
displayName: 'In-Memory Vector Store',
|
||||
name: 'vectorStoreInMemory',
|
||||
|
|
|
@ -213,7 +213,7 @@ class ExtendedPGVectorStore extends PGVectorStore {
|
|||
}
|
||||
}
|
||||
|
||||
export class VectorStorePGVector extends createVectorStoreNode({
|
||||
export class VectorStorePGVector extends createVectorStoreNode<ExtendedPGVectorStore>({
|
||||
meta: {
|
||||
description: 'Work with your data in Postgresql with the PGVector extension',
|
||||
icon: 'file:postgres.svg',
|
||||
|
@ -274,6 +274,7 @@ export class VectorStorePGVector extends createVectorStoreNode({
|
|||
|
||||
return await ExtendedPGVectorStore.initialize(embeddings, config);
|
||||
},
|
||||
|
||||
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
|
||||
// shared fields, because you need that strategy when creating the index.
|
||||
|
@ -307,6 +308,11 @@ export class VectorStorePGVector extends createVectorStoreNode({
|
|||
metadataColumnName: 'metadata',
|
||||
}) as ColumnOptions;
|
||||
|
||||
await PGVectorStore.fromDocuments(documents, embeddings, config);
|
||||
const vectorStore = await PGVectorStore.fromDocuments(documents, embeddings, config);
|
||||
vectorStore.client?.release();
|
||||
},
|
||||
|
||||
releaseVectorStoreClient(vectorStore) {
|
||||
vectorStore.client?.release();
|
||||
},
|
||||
}) {}
|
||||
|
|
|
@ -51,7 +51,7 @@ const insertFields: INodeProperties[] = [
|
|||
},
|
||||
];
|
||||
|
||||
export class VectorStorePinecone extends createVectorStoreNode({
|
||||
export class VectorStorePinecone extends createVectorStoreNode<PineconeStore>({
|
||||
meta: {
|
||||
displayName: 'Pinecone Vector Store',
|
||||
name: 'vectorStorePinecone',
|
||||
|
|
|
@ -79,7 +79,7 @@ const retrieveFields: INodeProperties[] = [
|
|||
},
|
||||
];
|
||||
|
||||
export class VectorStoreQdrant extends createVectorStoreNode({
|
||||
export class VectorStoreQdrant extends createVectorStoreNode<ExtendedQdrantVectorStore>({
|
||||
meta: {
|
||||
displayName: 'Qdrant Vector Store',
|
||||
name: 'vectorStoreQdrant',
|
||||
|
|
|
@ -41,7 +41,7 @@ const retrieveFields: INodeProperties[] = [
|
|||
|
||||
const updateFields: INodeProperties[] = [...insertFields];
|
||||
|
||||
export class VectorStoreSupabase extends createVectorStoreNode({
|
||||
export class VectorStoreSupabase extends createVectorStoreNode<SupabaseVectorStore>({
|
||||
meta: {
|
||||
description: 'Work with your data in Supabase Vector Store',
|
||||
icon: 'file:supabase.svg',
|
||||
|
|
|
@ -46,7 +46,7 @@ const retrieveFields: INodeProperties[] = [
|
|||
},
|
||||
];
|
||||
|
||||
export class VectorStoreZep extends createVectorStoreNode({
|
||||
export class VectorStoreZep extends createVectorStoreNode<ZepVectorStore | ZepCloudVectorStore>({
|
||||
meta: {
|
||||
displayName: 'Zep Vector Store',
|
||||
name: 'vectorStoreZep',
|
||||
|
|
|
@ -49,7 +49,7 @@ interface NodeMeta {
|
|||
operationModes?: NodeOperationMode[];
|
||||
}
|
||||
|
||||
export interface VectorStoreNodeConstructorArgs {
|
||||
export interface VectorStoreNodeConstructorArgs<T extends VectorStore = VectorStore> {
|
||||
meta: NodeMeta;
|
||||
methods?: {
|
||||
listSearch?: {
|
||||
|
@ -77,7 +77,8 @@ export interface VectorStoreNodeConstructorArgs {
|
|||
filter: Record<string, never> | undefined,
|
||||
embeddings: Embeddings,
|
||||
itemIndex: number,
|
||||
) => Promise<VectorStore>;
|
||||
) => Promise<T>;
|
||||
releaseVectorStoreClient?: (vectorStore: T) => void;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
function getOperationModeOptions(args: VectorStoreNodeConstructorArgs): INodePropertyOptions[] {
|
||||
function getOperationModeOptions<T extends VectorStore>(
|
||||
args: VectorStoreNodeConstructorArgs<T>,
|
||||
): INodePropertyOptions[] {
|
||||
const enabledOperationModes = args.meta.operationModes ?? DEFAULT_OPERATION_MODES;
|
||||
|
||||
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 {
|
||||
description: INodeTypeDescription = {
|
||||
displayName: args.meta.displayName,
|
||||
|
@ -334,38 +341,42 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
|
|||
embeddings,
|
||||
itemIndex,
|
||||
);
|
||||
const prompt = this.getNodeParameter('prompt', itemIndex) as string;
|
||||
const topK = this.getNodeParameter('topK', itemIndex, 4) as number;
|
||||
try {
|
||||
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 embeddedPrompt = await embeddings.embedQuery(prompt);
|
||||
const docs = await vectorStore.similaritySearchVectorWithScore(
|
||||
embeddedPrompt,
|
||||
topK,
|
||||
filter,
|
||||
);
|
||||
|
||||
const includeDocumentMetadata = this.getNodeParameter(
|
||||
'includeDocumentMetadata',
|
||||
itemIndex,
|
||||
true,
|
||||
) as boolean;
|
||||
const includeDocumentMetadata = this.getNodeParameter(
|
||||
'includeDocumentMetadata',
|
||||
itemIndex,
|
||||
true,
|
||||
) as boolean;
|
||||
|
||||
const serializedDocs = docs.map(([doc, score]) => {
|
||||
const document = {
|
||||
pageContent: doc.pageContent,
|
||||
...(includeDocumentMetadata ? { metadata: doc.metadata } : {}),
|
||||
};
|
||||
const serializedDocs = docs.map(([doc, score]) => {
|
||||
const document = {
|
||||
pageContent: doc.pageContent,
|
||||
...(includeDocumentMetadata ? { metadata: doc.metadata } : {}),
|
||||
};
|
||||
|
||||
return {
|
||||
json: { document, score },
|
||||
pairedItem: {
|
||||
item: itemIndex,
|
||||
},
|
||||
};
|
||||
});
|
||||
return {
|
||||
json: { document, score },
|
||||
pairedItem: {
|
||||
item: itemIndex,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
resultData.push(...serializedDocs);
|
||||
logAiEvent(this, 'ai-vector-store-searched', { query: prompt });
|
||||
resultData.push(...serializedDocs);
|
||||
logAiEvent(this, 'ai-vector-store-searched', { query: prompt });
|
||||
} finally {
|
||||
args.releaseVectorStoreClient?.(vectorStore);
|
||||
}
|
||||
}
|
||||
|
||||
return [resultData];
|
||||
|
@ -427,24 +438,28 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
|
|||
itemIndex,
|
||||
);
|
||||
|
||||
const { processedDocuments, serializedDocuments } = await processDocument(
|
||||
loader,
|
||||
itemData,
|
||||
itemIndex,
|
||||
);
|
||||
try {
|
||||
const { processedDocuments, serializedDocuments } = await processDocument(
|
||||
loader,
|
||||
itemData,
|
||||
itemIndex,
|
||||
);
|
||||
|
||||
if (processedDocuments?.length !== 1) {
|
||||
throw new NodeOperationError(this.getNode(), 'Single document per item expected');
|
||||
if (processedDocuments?.length !== 1) {
|
||||
throw new NodeOperationError(this.getNode(), 'Single document per item expected');
|
||||
}
|
||||
|
||||
resultData.push(...serializedDocuments);
|
||||
|
||||
// Use ids option to upsert instead of insert
|
||||
await vectorStore.addDocuments(processedDocuments, {
|
||||
ids: [documentId],
|
||||
});
|
||||
|
||||
logAiEvent(this, 'ai-vector-store-updated');
|
||||
} finally {
|
||||
args.releaseVectorStoreClient?.(vectorStore);
|
||||
}
|
||||
|
||||
resultData.push(...serializedDocuments);
|
||||
|
||||
// Use ids option to upsert instead of insert
|
||||
await vectorStore.addDocuments(processedDocuments, {
|
||||
ids: [documentId],
|
||||
});
|
||||
|
||||
logAiEvent(this, 'ai-vector-store-updated');
|
||||
}
|
||||
|
||||
return [resultData];
|
||||
|
@ -468,6 +483,9 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
|
|||
const vectorStore = await args.getVectorStoreClient(this, filter, embeddings, itemIndex);
|
||||
return {
|
||||
response: logWrapper(vectorStore, this),
|
||||
closeFunction: async () => {
|
||||
args.releaseVectorStoreClient?.(vectorStore);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
|
@ -491,23 +509,28 @@ export const createVectorStoreNode = (args: VectorStoreNodeConstructorArgs) =>
|
|||
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);
|
||||
|
||||
try {
|
||||
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);
|
||||
} finally {
|
||||
args.releaseVectorStoreClient?.(vectorStore);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
|
|
Loading…
Reference in a new issue