(TS) Fix PGVector implementation, where vector distance was inverted. (#4944)

2026-04-28 02:09:48 +07:00
parent 30ce028a71
commit ece7ff6b84
2 changed files with 129 additions and 1 deletions
@@ -251,7 +251,7 @@ export class PGVector implements VectorStore {
    return result.rows.map((row) => ({
      id: row.id,
      payload: row.payload,
-      score: row.distance,
+      score: Math.max(0, Math.min(1, 1 - Number(row.distance))),
    }));
  }

@@ -0,0 +1,128 @@
+/// <reference types="jest" />
+
+const searchRows = [
+  {
+    id: "a",
+    payload: { data: "exactly x-axis" },
+    distance: "0",
+  },
+  {
+    id: "b",
+    payload: { data: "close to x-axis" },
+    distance: "0.006116251198662548",
+  },
+  {
+    id: "c",
+    payload: { data: "y-axis" },
+    distance: "1",
+  },
+  {
+    id: "d",
+    payload: { data: "opposite x-axis" },
+    distance: "2",
+  },
+];
+
+function mockPgQuery(sql: string) {
+  if (sql.includes("SELECT 1 FROM pg_database")) {
+    return { rows: [{ "?column?": 1 }] };
+  }
+
+  if (sql.includes("FROM information_schema.tables")) {
+    return { rows: [{ table_name: "memories" }] };
+  }
+
+  if (sql.includes("vector <=> $1::vector AS distance")) {
+    return { rows: searchRows };
+  }
+
+  return { rows: [] };
+}
+
+jest.mock("pg", () => {
+  const clients: any[] = [];
+
+  const Client = jest.fn().mockImplementation((config: any) => {
+    const client = {
+      config,
+      connect: jest.fn().mockResolvedValue(undefined),
+      end: jest.fn().mockResolvedValue(undefined),
+      query: jest
+        .fn()
+        .mockImplementation(async (sql: string) => mockPgQuery(sql)),
+    };
+
+    clients.push(client);
+    return client;
+  });
+
+  return {
+    __esModule: true,
+    default: { Client },
+    Client,
+    __mock: { Client, clients },
+  };
+});
+
+import { PGVector } from "../src/vector_stores/pgvector";
+
+describe("PGVector - search()", () => {
+  beforeEach(() => {
+    const pg = require("pg");
+    pg.__mock.Client.mockClear();
+    pg.__mock.clients.length = 0;
+  });
+
+  test("converts cosine distance into a clamped similarity score", async () => {
+    const store = new PGVector({
+      collectionName: "memories",
+      user: "postgres",
+      password: "postgres",
+      host: "localhost",
+      port: 5432,
+      embeddingModelDims: 3,
+      dimension: 3,
+    } as any);
+
+    await store.initialize();
+
+    const results = await store.search([1, 0, 0], 4);
+
+    expect(results).toEqual([
+      {
+        id: "a",
+        payload: { data: "exactly x-axis" },
+        score: 1,
+      },
+      {
+        id: "b",
+        payload: { data: "close to x-axis" },
+        score: 0.9938837488013375,
+      },
+      {
+        id: "c",
+        payload: { data: "y-axis" },
+        score: 0,
+      },
+      {
+        id: "d",
+        payload: { data: "opposite x-axis" },
+        score: 0,
+      },
+    ]);
+
+    const pg = require("pg");
+    expect(pg.__mock.Client).toHaveBeenCalledTimes(2);
+
+    const activeClient = pg.__mock.clients[1];
+    expect(activeClient.query).toHaveBeenCalledWith(
+      expect.stringContaining("vector <=> $1::vector AS distance"),
+      ["[1,0,0]", 4],
+    );
+
+    for (const result of results) {
+      expect(result.score).toBeGreaterThanOrEqual(0);
+      expect(result.score).toBeLessThanOrEqual(1);
+    }
+  });
+});