From 5fd01d284a35a044602350d1faabaa02ee27c00d Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Fri, 25 Sep 2026 17:07:11 +0530 Subject: [PATCH] fix(ts-oss/turbopuffer): bound euclidean_squared distance to a similarity score (#6580) --- .../src/oss/src/vector_stores/turbopuffer.ts | 13 ++- .../src/oss/tests/turbopuffer.score.test.ts | 80 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 mem0-ts/src/oss/tests/turbopuffer.score.test.ts diff --git a/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts b/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts index 8979818d2..30e032630 100644 --- a/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts +++ b/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts @@ -305,7 +305,18 @@ export class TurbopufferDB implements VectorStore { private parseRows(rows: any[]): VectorStoreResult[] { return rows.map((row) => { const { id, $dist, vector, ...rest } = row; - const score = $dist != null ? 1 - $dist : undefined; + let score: number | undefined; + if ($dist == null) { + score = undefined; + } else if (this.distanceMetric === "euclidean_squared") { + // euclidean_squared $dist is an unbounded squared distance, so 1 - $dist + // goes negative for any $dist > 1 and inverts ranking. Convert it to a + // bounded higher-is-better score, mirroring the milvus/baidu stores. + score = 1 / (1 + $dist); + } else { + // Cosine distance is in [0, 2]; 1 - $dist stays a meaningful similarity. + score = 1 - $dist; + } return { id: String(id), payload: rest, score }; }); } diff --git a/mem0-ts/src/oss/tests/turbopuffer.score.test.ts b/mem0-ts/src/oss/tests/turbopuffer.score.test.ts new file mode 100644 index 000000000..b6e8ba844 --- /dev/null +++ b/mem0-ts/src/oss/tests/turbopuffer.score.test.ts @@ -0,0 +1,80 @@ +/// +/** + * Turbopuffer vector store — score conversion unit tests. + * + * Drives parseRows() through the public search() API with a virtually-mocked + * @turbopuffer/turbopuffer peer, asserting the score returned per metric. + */ + +const mockQuery = jest.fn(); + +jest.mock( + "@turbopuffer/turbopuffer", + () => ({ + __esModule: true, + default: class { + namespace() { + return { query: mockQuery }; + } + }, + }), + { virtual: true }, +); + +import { TurbopufferDB } from "../src/vector_stores/turbopuffer"; + +function makeStore(distanceMetric?: string) { + return new TurbopufferDB({ + apiKey: "test-key", + collectionName: "mem0", + ...(distanceMetric ? { distanceMetric } : {}), + } as any); +} + +async function scoreFor( + distanceMetric: string | undefined, + row: Record, +): Promise { + mockQuery.mockResolvedValueOnce({ rows: [row] }); + const results = await makeStore(distanceMetric).search([0.1, 0.2, 0.3], 5); + return results[0].score; +} + +describe("TurbopufferDB score conversion", () => { + it("keeps cosine distance as 1 - dist (default metric)", async () => { + // cosine_distance is the default; a distance of 0.25 -> similarity 0.75. + expect(await scoreFor(undefined, { id: "a", $dist: 0.25 })).toBeCloseTo( + 0.75, + 10, + ); + }); + + it("bounds an unbounded euclidean_squared distance to a 0..1 similarity", async () => { + // $dist = 4.0 is a squared distance. 1 - 4.0 = -3.0 would invert ranking; + // 1 / (1 + 4.0) = 0.2 keeps it higher-is-better and in range. + const score = await scoreFor("euclidean_squared", { id: "a", $dist: 4.0 }); + expect(score).toBeCloseTo(0.2, 10); + expect(score!).toBeGreaterThanOrEqual(0); + expect(score!).toBeLessThanOrEqual(1); + }); + + it("ranks a nearer euclidean_squared hit above a farther one", async () => { + mockQuery.mockResolvedValueOnce({ + rows: [ + { id: "near", $dist: 1.0 }, + { id: "far", $dist: 9.0 }, + ], + }); + const results = await makeStore("euclidean_squared").search( + [0.1, 0.2, 0.3], + 5, + ); + const near = results.find((r) => r.id === "near")!; + const far = results.find((r) => r.id === "far")!; + expect(near.score!).toBeGreaterThan(far.score!); + }); + + it("preserves an undefined score when the row has no distance", async () => { + expect(await scoreFor("euclidean_squared", { id: "a" })).toBeUndefined(); + }); +});