mudhorn/n8n
1
0
Fork 0
mirror of https://github.com/n8n-io/n8n.git synced 2025-03-05 20:50:17 -08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix(Postgres PGVector Store Node): Fix filtering in retriever mode (#11075)

Browse source
This commit is contained in:
oleg 2024-10-03 12:50:35 +02:00 committed by GitHub
parent 2596ddbe8d
commit dbd2ae1995
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 40 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
packages/@n8n/nodes-langchain/nodes/vector_store/VectorStorePGVector/VectorStorePGVector.node.ts
View file

@ -1,14 +1,16 @@
import { type INodeProperties } from 'n8n-workflow';
import { import {
PGVectorStore, PGVectorStore,
type DistanceStrategy, type DistanceStrategy,
type PGVectorStoreArgs, type PGVectorStoreArgs,
} from '@langchain/community/vectorstores/pgvector'; } 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 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 type pg from 'pg';
import { createVectorStoreNode } from '../shared/createVectorStoreNode';
import { metadataFilterField } from '../../../utils/sharedFields'; import { metadataFilterField } from '../../../utils/sharedFields';
import { createVectorStoreNode } from '../shared/createVectorStoreNode';
type CollectionOptions = { type CollectionOptions = {
useCollection?: boolean; 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<ExtendedPGVectorStore> {
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({ export const VectorStorePGVector = createVectorStoreNode({
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',
displayName: 'Postgres PGVector Store', displayName: 'Postgres PGVector Store',
docsUrl: 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', name: 'vectorStorePGVector',
credentials: [ credentials: [
{ {
@ -236,7 +271,7 @@ export const VectorStorePGVector = createVectorStoreNode({
'cosine', 'cosine',
) as DistanceStrategy; ) as DistanceStrategy;
return await PGVectorStore.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