Bugfix/create postgresql table if not exists for vectorstore (#5573)

create postgresql table if not exists for vectorstore

Author: HenryHengZJ

packages/components/nodes/vectorstores/Postgres/driver/Base.ts

@@ -2,7 +2,7 @@ import { VectorStore } from '@langchain/core/vectorstores'
 import { getCredentialData, getCredentialParam, ICommonObject, INodeData } from '../../../../src'
 import { Document } from '@langchain/core/documents'
 import { Embeddings } from '@langchain/core/embeddings'
-import { getDatabase, getHost, getPort, getSSL, getTableName } from '../utils'
+import { getDatabase, getHost, getPort, getSchemaName, getSSL, getTableName } from '../utils'
 
 export abstract class VectorStoreDriver {
     constructor(protected nodeData: INodeData, protected options: ICommonObject) {}
@@ -35,6 +35,18 @@ export abstract class VectorStoreDriver {
         return this.sanitizeTableName(getTableName(this.nodeData))
     }
 
+    getSchemaName() {
+        const schemaName = getSchemaName(this.nodeData)
+        return schemaName ? this.sanitizeTableName(schemaName) : undefined
+    }
+
+    getTablePath() {
+        const schemaName = this.getSchemaName()
+        const tableName = this.getTableName()
+        if (!schemaName) return `"${tableName}"`
+        return `"${schemaName}"."${tableName}"`
+    }
+
     getEmbeddings() {
         return this.nodeData.inputs?.embeddings as Embeddings
     }

packages/components/nodes/vectorstores/Postgres/driver/TypeORM.ts

@@ -52,16 +52,25 @@ export class TypeORMDriver extends VectorStoreDriver {
     async getArgs(): Promise<TypeORMVectorStoreArgs> {
         return {
             postgresConnectionOptions: await this.getPostgresConnectionOptions(),
-            tableName: this.getTableName()
+            tableName: this.getTableName(),
+            schemaName: this.getSchemaName()
         }
     }
 
     async instanciate(metadataFilters?: any) {
-        return this.adaptInstance(await TypeORMVectorStore.fromDataSource(this.getEmbeddings(), await this.getArgs()), metadataFilters)
+        return this.adaptInstance(
+            await TypeORMVectorStore.fromDataSource(this.getEmbeddings(), await this.getArgs()),
+            metadataFilters,
+            this.getTablePath()
+        )
     }
 
     async fromDocuments(documents: Document[]) {
-        return this.adaptInstance(await TypeORMVectorStore.fromDocuments(documents, this.getEmbeddings(), await this.getArgs()))
+        return this.adaptInstance(
+            await TypeORMVectorStore.fromDocuments(documents, this.getEmbeddings(), await this.getArgs()),
+            undefined,
+            this.getTablePath()
+        )
     }
 
     sanitizeDocuments(documents: Document[]) {
@@ -73,8 +82,8 @@ export class TypeORMDriver extends VectorStoreDriver {
         return documents
     }
 
-    protected async adaptInstance(instance: TypeORMVectorStore, metadataFilters?: any): Promise<VectorStore> {
-        const tableName = this.getTableName()
+    protected async adaptInstance(instance: TypeORMVectorStore, metadataFilters?: any, tablePath?: string): Promise<VectorStore> {
+        const effectiveTablePath = tablePath ?? this.getTablePath()
 
         // Rewrite the method to use pg pool connection instead of the default connection
         /* Otherwise a connection error is displayed when the chain tries to execute the function
@@ -86,7 +95,7 @@ export class TypeORMDriver extends VectorStoreDriver {
             return await TypeORMDriver.similaritySearchVectorWithScore(
                 query,
                 k,
-                tableName,
+                effectiveTablePath,
                 await this.getPostgresConnectionOptions(),
                 filter ?? metadataFilters,
                 this.computedOperatorString
@@ -141,6 +150,8 @@ export class TypeORMDriver extends VectorStoreDriver {
 
         instance.addDocuments = async (documents: Document[], options?: { ids?: string[] }): Promise<void> => {
             const texts = documents.map(({ pageContent }) => pageContent)
+            // Ensure table exists before adding documents (this will create the table if it does not exist)
+            await this.ensureTableInDatabase(instance, effectiveTablePath)
             return (instance.addVectors as any)(await this.getEmbeddings().embedDocuments(texts), documents, options)
         }
 
@@ -162,10 +173,26 @@ export class TypeORMDriver extends VectorStoreDriver {
         }
     }
 
+    /**
+     * Ensures the table exists in the database with the correct schema.
+     * Creates the pgvector extension and table if they don't exist.
+     */
+    async ensureTableInDatabase(instance: TypeORMVectorStore, tablePath: string): Promise<void> {
+        await instance.appDataSource.query('CREATE EXTENSION IF NOT EXISTS vector;')
+        await instance.appDataSource.query(`
+            CREATE TABLE IF NOT EXISTS ${tablePath} (
+                "id" uuid NOT NULL DEFAULT gen_random_uuid() PRIMARY KEY,
+                "pageContent" text,
+                metadata jsonb,
+                embedding vector
+            );
+        `)
+    }
+
     static similaritySearchVectorWithScore = async (
         query: number[],
         k: number,
-        tableName: string,
+        tablePath: string,
         postgresConnectionOptions: ICommonObject,
         filter?: any,
         distanceOperator: string = '<=>'
@@ -186,7 +213,7 @@ export class TypeORMDriver extends VectorStoreDriver {
 
         const queryString = `
             SELECT *, embedding ${distanceOperator} $1 as "_distance"
-            FROM ${tableName}
+            FROM ${tablePath}
             WHERE ((metadata @> $2) AND NOT (metadata ? '${FLOWISE_CHATID}')) ${chatflowOr}
             ORDER BY "_distance" ASC
             LIMIT $3;`

packages/components/nodes/vectorstores/Postgres/utils.ts

@@ -23,3 +23,7 @@ export function getTableName(nodeData?: INodeData) {
 export function getContentColumnName(nodeData?: INodeData) {
     return defaultChain(nodeData?.inputs?.contentColumnName, process.env.POSTGRES_VECTORSTORE_CONTENT_COLUMN_NAME, 'pageContent')
 }
+
+export function getSchemaName(nodeData?: INodeData) {
+    return defaultChain(nodeData?.inputs?.schemaName, process.env.POSTGRES_VECTORSTORE_SCHEMA_NAME)
+}