From dbd2ae199506a24c2df4c983111a56f2adf63eee Mon Sep 17 00:00:00 2001 From: oleg Date: Thu, 3 Oct 2024 12:50:35 +0200 Subject: [PATCH] fix(Postgres PGVector Store Node): Fix filtering in retriever mode (#11075) --- .../VectorStorePGVector.node.ts | 45 ++++++++++++++++--- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts index 7e14eb488..0c9a148be 100644 --- a/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts +++ b/packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts @@ -1,14 +1,16 @@ -import { type INodeProperties } from 'n8n-workflow'; import { PGVectorStore, type DistanceStrategy, type PGVectorStoreArgs, } from '@langchain/community/vectorstores/pgvector'; -import { configurePostgres } from 'n8n-nodes-base/dist/nodes/Postgres/v2/transport'; +import type { EmbeddingsInterface } from '@langchain/core/embeddings'; import type { PostgresNodeCredentials } from 'n8n-nodes-base/dist/nodes/Postgres/v2/helpers/interfaces'; +import { configurePostgres } from 'n8n-nodes-base/dist/nodes/Postgres/v2/transport'; +import type { INodeProperties } from 'n8n-workflow'; import type pg from 'pg'; -import { createVectorStoreNode } from '../shared/createVectorStoreNode'; + import { metadataFilterField } from '../../../utils/sharedFields'; +import { createVectorStoreNode } from '../shared/createVectorStoreNode'; type CollectionOptions = { useCollection?: boolean; @@ -177,13 +179,46 @@ const retrieveFields: INodeProperties[] = [ }, ]; +/** + * Extended PGVectorStore class to handle custom filtering. + * This wrapper is necessary because when used as a retriever, + * similaritySearchVectorWithScore should use this.filter instead of + * expecting it from the parameter + */ +class ExtendedPGVectorStore extends PGVectorStore { + static async initialize( + embeddings: EmbeddingsInterface, + args: PGVectorStoreArgs & { dimensions?: number }, + ): Promise { + const { dimensions, ...rest } = args; + const postgresqlVectorStore = new this(embeddings, rest); + + await postgresqlVectorStore._initializeClient(); + await postgresqlVectorStore.ensureTableInDatabase(dimensions); + if (postgresqlVectorStore.collectionTableName) { + await postgresqlVectorStore.ensureCollectionTableInDatabase(); + } + + return postgresqlVectorStore; + } + + async similaritySearchVectorWithScore( + query: number[], + k: number, + filter?: PGVectorStore['FilterType'], + ) { + const mergedFilter = { ...this.filter, ...filter }; + return await super.similaritySearchVectorWithScore(query, k, mergedFilter); + } +} + export const VectorStorePGVector = createVectorStoreNode({ meta: { description: 'Work with your data in Postgresql with the PGVector extension', icon: 'file:postgres.svg', displayName: 'Postgres PGVector Store', docsUrl: - 'https://docs.n8n.io/integrations/builtin/cluster-nodes/root-nodes/n8n-nodes-langchain.vectorstoresupabase/', + 'https://docs.n8n.io/integrations/builtin/cluster-nodes/root-nodes/n8n-nodes-langchain.vectorstorepgvector/', name: 'vectorStorePGVector', credentials: [ { @@ -236,7 +271,7 @@ export const VectorStorePGVector = createVectorStoreNode({ 'cosine', ) as DistanceStrategy; - return await PGVectorStore.initialize(embeddings, config); + 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