community[patch]: Adds support for passing schemaName to pgvector

libs/langchain-community/src/vectorstores/pgvector.ts

@@ -17,6 +17,8 @@ export interface PGVectorStoreArgs {
   collectionTableName?: string;
   collectionName?: string;
   collectionMetadata?: Metadata | null;
+  schemaName?: string | null;
+  extensionSchemaName?: String | null;
   columns?: {
     idColumnName?: string;
     vectorColumnName?: string;
@@ -51,12 +53,16 @@ export class PGVectorStore extends VectorStore {
 
   collectionMetadata: Metadata | null;
 
+  schemaName: string | null;
+
   idColumnName: string;
 
   vectorColumnName: string;
 
   contentColumnName: string;
 
+  extensionSchemaName: string | null;
+
   metadataColumnName: string;
 
   filter?: Metadata;
@@ -82,6 +88,9 @@ export class PGVectorStore extends VectorStore {
     this.collectionTableName = config.collectionTableName;
     this.collectionName = config.collectionName ?? "langchain";
     this.collectionMetadata = config.collectionMetadata ?? null;
+    this.schemaName = config.schemaName ?? null;
+    this.extensionSchemaName = config.extensionSchemaName ?? null;
+
     this.filter = config.filter;
 
     this.vectorColumnName = config.columns?.vectorColumnName ?? "embedding";
@@ -237,11 +246,16 @@ export class PGVectorStore extends VectorStore {
       .map((_, j) => this.generatePlaceholderForRowAt(j, columns.length))
       .join(", ");
 
-    const text = `
+    const text = this.schemaName == null ? `
       INSERT INTO ${this.tableName}(
         ${columns.map((column) => `"${column}"`).join(", ")}
       )
       VALUES ${valuesPlaceholders}
+    ` : `
+      INSERT INTO ${this.schemaName}.${this.tableName}(
+        ${columns.map((column) => `"${column}"`).join(", ")}
+      )
+      VALUES ${valuesPlaceholders}
     `;
     return text;
   }
@@ -322,11 +336,16 @@ export class PGVectorStore extends VectorStore {
     // Set parameters of dynamically generated query
     const params = collectionId ? [ids, collectionId] : [ids];
 
-    const queryString = `
+    const queryString = this.schemaName == null ? `
       DELETE FROM ${this.tableName}
       WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
       this.idColumnName
     } = ANY($1::uuid[])
+    ` : `
+      DELETE FROM ${this.schemaName}.${this.tableName}
+      WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
+      this.idColumnName
+    } = ANY($1::uuid[])
     `;
     await this.pool.query(queryString, params);
   }
@@ -347,11 +366,16 @@ export class PGVectorStore extends VectorStore {
     // Set parameters of dynamically generated query
     const params = collectionId ? [filter, collectionId] : [filter];
 
-    const queryString = `
+    const queryString = this.schemaName == null ? `
       DELETE FROM ${this.tableName}
       WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
       this.metadataColumnName
     }::jsonb @> $1
+    ` : `
+      DELETE FROM ${this.schemaName}.${this.tableName}
+      WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
+      this.metadataColumnName
+    }::jsonb @> $1
     `;
     return await this.pool.query(queryString, params);
   }
@@ -454,13 +478,20 @@ export class PGVectorStore extends VectorStore {
       ? `WHERE ${whereClauses.join(" AND ")}`
       : "";
 
-    const queryString = `
+    const queryString = this.schemaName == null ? `
       SELECT *, ${this.vectorColumnName} <=> $1 as "_distance"
       FROM ${this.tableName}
       ${whereClause}
       ${collectionId ? "AND collection_id = $3" : ""}
       ORDER BY "_distance" ASC
       LIMIT $2;
+      ` : `
+      SELECT *, ${this.vectorColumnName} <=> $1 as "_distance"
+      FROM ${this.schemaName}.${this.tableName}
+      ${whereClause}
+      ${collectionId ? "AND collection_id = $3" : ""}
+      ORDER BY "_distance" ASC
+      LIMIT $2;
       `;
 
     const documents = (await this.pool.query(queryString, parameters)).rows;
@@ -485,17 +516,26 @@ export class PGVectorStore extends VectorStore {
    * @returns Promise that resolves when the table has been ensured.
    */
   async ensureTableInDatabase(): Promise<void> {
-    await this.pool.query("CREATE EXTENSION IF NOT EXISTS vector;");
-    await this.pool.query('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";');
-
-    await this.pool.query(`
+    const vectorQuery = this.extensionSchemaName == null ? "CREATE EXTENSION IF NOT EXISTS vector;" : `CREATE EXTENSION IF NOT EXISTS vector WITH SCHEMA ${this.extensionSchemaName};`
+    const uuidQuery = this.extensionSchemaName == null ? 'CREATE EXTENSION IF NOT EXISTS "uuid-ossp";' : `CREATE EXTENSION IF NOT EXISTS "uuid-ossp" WITH SCHEMA ${this.extensionSchemaName};`
+     const tableQuery = this.schemaName == null ? `
       CREATE TABLE IF NOT EXISTS ${this.tableName} (
         "${this.idColumnName}" uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
         "${this.contentColumnName}" text,
         "${this.metadataColumnName}" jsonb,
         "${this.vectorColumnName}" vector
       );
-    `);
+    ` : `
+      CREATE TABLE IF NOT EXISTS ${this.schemaName}.${this.tableName} (
+        "${this.idColumnName}" uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
+        "${this.contentColumnName}" text,
+        "${this.metadataColumnName}" jsonb,
+        "${this.vectorColumnName}" ${this.schemaName}.vector
+      );
+    `
+    await this.pool.query(vectorQuery);
+    await this.pool.query(uuidQuery);
+    await this.pool.query(tableQuery);
   }
 
   /**
@@ -506,7 +546,7 @@ export class PGVectorStore extends VectorStore {
    */
   async ensureCollectionTableInDatabase(): Promise<void> {
     try {
-      await this.pool.query(`
+      const queryString = this.schemaName == null ? `
         CREATE TABLE IF NOT EXISTS ${this.collectionTableName} (
           uuid uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
           name character varying,
@@ -521,7 +561,23 @@ export class PGVectorStore extends VectorStore {
           FOREIGN KEY (collection_id)
           REFERENCES ${this.collectionTableName}(uuid)
           ON DELETE CASCADE;
-      `);
+      ` : `
+        CREATE TABLE IF NOT EXISTS ${this.schemaName}.${this.collectionTableName} (
+          uuid uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
+          name character varying,
+          cmetadata jsonb
+        );
+
+        ALTER TABLE ${this.schemaName}.${this.tableName}
+          ADD COLUMN collection_id uuid;
+
+        ALTER TABLE ${this.schemaName}.${this.tableName}
+          ADD CONSTRAINT ${this.schemaName}.${this.tableName}_collection_id_fkey
+          FOREIGN KEY (collection_id)
+          REFERENCES ${this.collectionTableName}(uuid)
+          ON DELETE CASCADE;
+      `
+      await this.pool.query(queryString);
     } catch (e) {
       if (!(e as Error).message.includes("already exists")) {
         console.error(e);