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

)

* Adds support for passing schemaName to pgvector

* Updates refrences to tableName to use schemaName where missing

* Adds missing var to pool call

* Adds support for nesting extensions in schema

* Adds support for nesting extensions under schemas

* Adds wrapping to handle capital letters

* fixes formatting

* Removes duplicate queries

* Removes extra string wrapping

* Fixes syntax error

* Removes string wrapper around tableName where schemaName isnt present

* Removes additional quote wrapping

* Adds computedTableName to class

* Uses tableName for constraint identifier

* yarn format

* Fixes edgecase where operator was out of scope due to order of operations with schema call

* Adds check to handle unused extensionSchemaName

* Refactors updating search_path to use operator function instead

* Change to getter method, revert inadvertent change

* Revert

---------

Co-authored-by: jacoblee93 <jacoblee93@gmail.com>

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";
@@ -98,6 +107,12 @@ export class PGVectorStore extends VectorStore {
       !!config.verbose;
   }
 
+  get computedTableName() {
+    return this.schemaName == null
+      ? `${this.tableName}`
+      : `"${this.schemaName}"."${this.tableName}"`;
+  }
+
   /**
    * Static method to create a new `PGVectorStore` instance from a
    * connection. It creates a table if one does not exist, and calls
@@ -238,7 +253,7 @@ export class PGVectorStore extends VectorStore {
       .join(", ");
 
     const text = `
-      INSERT INTO ${this.tableName}(
+      INSERT INTO ${this.computedTableName}(
         ${columns.map((column) => `"${column}"`).join(", ")}
       )
       VALUES ${valuesPlaceholders}
@@ -323,7 +338,7 @@ export class PGVectorStore extends VectorStore {
     const params = collectionId ? [ids, collectionId] : [ids];
 
     const queryString = `
-      DELETE FROM ${this.tableName}
+      DELETE FROM ${this.computedTableName}
       WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
       this.idColumnName
     } = ANY($1::uuid[])
@@ -348,7 +363,7 @@ export class PGVectorStore extends VectorStore {
     const params = collectionId ? [filter, collectionId] : [filter];
 
     const queryString = `
-      DELETE FROM ${this.tableName}
+      DELETE FROM ${this.computedTableName}
       WHERE ${collectionId ? "collection_id = $2 AND " : ""}${
       this.metadataColumnName
     }::jsonb @> $1
@@ -454,9 +469,14 @@ export class PGVectorStore extends VectorStore {
       ? `WHERE ${whereClauses.join(" AND ")}`
       : "";
 
+    const operatorString =
+      this.extensionSchemaName !== null
+        ? `OPERATOR(${this.extensionSchemaName}.<=>)`
+        : "<=>";
+
     const queryString = `
-      SELECT *, ${this.vectorColumnName} <=> $1 as "_distance"
-      FROM ${this.tableName}
+      SELECT *, "${this.vectorColumnName}" ${operatorString} $1 as "_distance"
+      FROM ${this.computedTableName}
       ${whereClause}
       ${collectionId ? "AND collection_id = $3" : ""}
       ORDER BY "_distance" ASC
@@ -485,17 +505,29 @@ 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(`
-      CREATE TABLE IF NOT EXISTS ${this.tableName} (
+    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 extensionName =
+      this.extensionSchemaName == null
+        ? "vector"
+        : `"${this.extensionSchemaName}"."vector"`;
+    const tableQuery = `
+      CREATE TABLE IF NOT EXISTS ${this.computedTableName} (
         "${this.idColumnName}" uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
         "${this.contentColumnName}" text,
         "${this.metadataColumnName}" jsonb,
-        "${this.vectorColumnName}" vector
+        "${this.vectorColumnName}" ${extensionName}
       );
-    `);
+    `;
+    await this.pool.query(vectorQuery);
+    await this.pool.query(uuidQuery);
+    await this.pool.query(tableQuery);
   }
 
   /**
@@ -506,22 +538,23 @@ export class PGVectorStore extends VectorStore {
    */
   async ensureCollectionTableInDatabase(): Promise<void> {
     try {
-      await this.pool.query(`
+      const queryString = `
         CREATE TABLE IF NOT EXISTS ${this.collectionTableName} (
           uuid uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
           name character varying,
           cmetadata jsonb
         );
 
-        ALTER TABLE ${this.tableName}
+        ALTER TABLE ${this.computedTableName}
           ADD COLUMN collection_id uuid;
 
-        ALTER TABLE ${this.tableName}
+        ALTER TABLE ${this.computedTableName}
           ADD CONSTRAINT ${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);