diff --git a/mem0-ts/src/oss/src/vector_stores/pgvector.ts b/mem0-ts/src/oss/src/vector_stores/pgvector.ts index 542527757..ac957b168 100644 --- a/mem0-ts/src/oss/src/vector_stores/pgvector.ts +++ b/mem0-ts/src/oss/src/vector_stores/pgvector.ts @@ -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))), })); } diff --git a/mem0-ts/src/oss/tests/pgvector.unit.test.ts b/mem0-ts/src/oss/tests/pgvector.unit.test.ts new file mode 100644 index 000000000..ef4827ced --- /dev/null +++ b/mem0-ts/src/oss/tests/pgvector.unit.test.ts @@ -0,0 +1,128 @@ +/// + +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); + } + }); +});