community[patch]: Add support for inner product and Euclidean distanc…

…e to `PGVector` (langchain-ai#4781) * Add support for inner product and euclidean distance to PGVector. * Fix formatting errors. * Rename SupportedVectorTypes to DistanceStrategy. Update PGVector docs.
lokesh-couchbase · Mar 15, 2024 · fdbc858 · fdbc858
1 parent 231e475
commit fdbc858
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 7 deletions.
diff --git a/examples/src/indexes/vector_stores/pgvector_vectorstore/pgvector.ts b/examples/src/indexes/vector_stores/pgvector_vectorstore/pgvector.ts
@@ -1,5 +1,8 @@
 import { OpenAIEmbeddings } from "@langchain/openai";
-import { PGVectorStore } from "@langchain/community/vectorstores/pgvector";
+import {
+  DistanceStrategy,
+  PGVectorStore,
+} from "@langchain/community/vectorstores/pgvector";
 import { PoolConfig } from "pg";
 
 // First, follow set-up instructions at
@@ -21,6 +24,8 @@ const config = {
     contentColumnName: "content",
     metadataColumnName: "metadata",
   },
+  // supported distance strategies: cosine (default), innerProduct, or euclidean
+  distanceStrategy: "cosine" as DistanceStrategy,
 };
 
 const pgvectorStore = await PGVectorStore.initialize(

diff --git a/libs/langchain-community/src/vectorstores/pgvector.ts b/libs/langchain-community/src/vectorstores/pgvector.ts
@@ -6,6 +6,8 @@ import { getEnvironmentVariable } from "@langchain/core/utils/env";
 
 type Metadata = Record<string, unknown>;
 
+export type DistanceStrategy = "cosine" | "innerProduct" | "euclidean";
+
 /**
  * Interface that defines the arguments required to create a
  * `PGVectorStore` instance. It includes Postgres connection options,
@@ -35,6 +37,7 @@ export interface PGVectorStoreArgs {
    */
   chunkSize?: number;
   ids?: string[];
+  distanceStrategy?: DistanceStrategy;
 }
 
 /**
@@ -76,6 +79,8 @@ export class PGVectorStore extends VectorStore {
 
   chunkSize = 500;
 
+  distanceStrategy?: DistanceStrategy = "cosine";
+
   _vectorstoreType(): string {
     return "pgvector";
   }
@@ -112,6 +117,7 @@ export class PGVectorStore extends VectorStore {
     const pool = config.pool ?? new pg.Pool(config.postgresConnectionOptions);
     this.pool = pool;
     this.chunkSize = config.chunkSize ?? 500;
+    this.distanceStrategy = config.distanceStrategy ?? this.distanceStrategy;
 
     this._verbose =
       getEnvironmentVariable("LANGCHAIN_VERBOSE") === "true" ??
@@ -130,6 +136,27 @@ export class PGVectorStore extends VectorStore {
       : `"${this.schemaName}"."${this.collectionTableName}"`;
   }
 
+  get computedOperatorString() {
+    let operator: string;
+    switch (this.distanceStrategy) {
+      case "cosine":
+        operator = "<=>";
+        break;
+      case "innerProduct":
+        operator = "<#>";
+        break;
+      case "euclidean":
+        operator = "<->";
+        break;
+      default:
+        throw new Error(`Unknown distance strategy: ${this.distanceStrategy}`);
+    }
+
+    return this.extensionSchemaName !== null
+      ? `OPERATOR(${this.extensionSchemaName}.${operator})`
+      : operator;
+  }
+
   /**
    * Static method to create a new `PGVectorStore` instance from a
    * connection. It creates a table if one does not exist, and calls
@@ -487,13 +514,8 @@ export class PGVectorStore extends VectorStore {
       ? `WHERE ${whereClauses.join(" AND ")}`
       : "";
 
-    const operatorString =
-      this.extensionSchemaName !== null
-        ? `OPERATOR(${this.extensionSchemaName}.<=>)`
-        : "<=>";
-
     const queryString = `
-      SELECT *, "${this.vectorColumnName}" ${operatorString} $1 as "_distance"
+      SELECT *, "${this.vectorColumnName}" ${this.computedOperatorString} $1 as "_distance"
       FROM ${this.computedTableName}
       ${whereClause}
       ORDER BY "_distance" ASC

diff --git a/libs/langchain-community/src/vectorstores/tests/pgvector/pgvector.int.test.ts b/libs/langchain-community/src/vectorstores/tests/pgvector/pgvector.int.test.ts
@@ -219,6 +219,37 @@ describe("PGVectorStore", () => {
       throw e;
     }
   });
+
+  test("PGvector supports different vector types", async () => {
+    // verify by asserting different pgvector operators based on vector type
+    pgvectorVectorStore.distanceStrategy = "cosine";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual("<=>");
+
+    pgvectorVectorStore.distanceStrategy = "innerProduct";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual("<#>");
+
+    pgvectorVectorStore.distanceStrategy = "euclidean";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual("<->");
+
+    // verify with extensionSchemaName
+    pgvectorVectorStore.distanceStrategy = "cosine";
+    pgvectorVectorStore.extensionSchemaName = "schema1";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual(
+      "OPERATOR(schema1.<=>)"
+    );
+
+    pgvectorVectorStore.distanceStrategy = "innerProduct";
+    pgvectorVectorStore.extensionSchemaName = "schema2";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual(
+      "OPERATOR(schema2.<#>)"
+    );
+
+    pgvectorVectorStore.distanceStrategy = "euclidean";
+    pgvectorVectorStore.extensionSchemaName = "schema3";
+    expect(pgvectorVectorStore.computedOperatorString).toEqual(
+      "OPERATOR(schema3.<->)"
+    );
+  });
 });
 
 describe.skip("PGVectorStore with collection", () => {