From 2cc060fd761f5a3c49a842e6da9b02f19975f0f7 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Mon, 6 Jul 2026 13:27:00 -0400 Subject: [PATCH] feat(vector-stores): add Turbopuffer provider to TypeScript OSS SDK (#5801) Co-authored-by: kartik-mem0 --- docs/components/vectordbs/dbs/turbopuffer.mdx | 56 ++- mem0-ts/package.json | 1 + mem0-ts/pnpm-lock.yaml | 16 + .../tests/vector_stores/turbopuffer.test.ts | 412 ++++++++++++++++++ mem0-ts/src/oss/src/utils/factory.ts | 3 + .../src/oss/src/vector_stores/turbopuffer.ts | 238 ++++++++++ mem0-ts/tsup.config.ts | 1 + 7 files changed, 725 insertions(+), 2 deletions(-) create mode 100644 mem0-ts/src/oss/src/tests/vector_stores/turbopuffer.test.ts create mode 100644 mem0-ts/src/oss/src/vector_stores/turbopuffer.ts diff --git a/docs/components/vectordbs/dbs/turbopuffer.mdx b/docs/components/vectordbs/dbs/turbopuffer.mdx index 502e69655..825d8f55c 100644 --- a/docs/components/vectordbs/dbs/turbopuffer.mdx +++ b/docs/components/vectordbs/dbs/turbopuffer.mdx @@ -6,7 +6,8 @@ description: "Use Turbopuffer as a serverless vector database in Mem0 for low-la ### Usage -```python + +```python Python import os from mem0 import Memory @@ -39,6 +40,36 @@ m.add(messages, user_id="alice", metadata={"category": "movies"}) results = m.search(query="sci-fi recommendations", filters={"user_id": "alice"}) ``` +```typescript TypeScript +import { Memory } from "mem0ai/oss"; + +// Set TURBOPUFFER_API_KEY in your environment, or pass it as config.apiKey below. +const config = { + vectorStore: { + provider: "turbopuffer", + config: { + collectionName: "movie_preferences", + region: "gcp-us-central1", + }, + }, +}; + +const memory = new Memory(config); + +const messages = [ + { role: "user", content: "I'm planning to watch a movie tonight. Any recommendations?" }, + { role: "assistant", content: "How about thriller movies? They can be quite engaging." }, + { role: "user", content: "I'm not a big fan of thrillers but I love sci-fi." }, + { role: "assistant", content: "Got it! I'll suggest sci-fi movies instead." }, +]; + +await memory.add(messages, { userId: "alice", metadata: { category: "movies" } }); + +// Search memories +const results = await memory.search("sci-fi recommendations", { userId: "alice" }); +``` + + ### Config Here are the parameters available for configuring Turbopuffer: @@ -53,6 +84,10 @@ Here are the parameters available for configuring Turbopuffer: | `batch_size` | Batch size for bulk operations | `100` | | `extra_params` | Additional parameters for the Turbopuffer client | `None` | + +**TypeScript (Node.js) config keys** are camelCase: `collectionName`, `apiKey`, `region`, `distanceMetric`, and `batchSize`. The TypeScript SDK infers the vector dimension from your embedder, so `embeddingModelDims` is not required. + + ### Regions | Region | Location | @@ -62,7 +97,8 @@ Here are the parameters available for configuring Turbopuffer: ### Config Example -```python + +```python Python config = { "vector_store": { "provider": "turbopuffer", @@ -77,3 +113,19 @@ config = { } } ``` + +```typescript TypeScript +const config = { + vectorStore: { + provider: "turbopuffer", + config: { + collectionName: "my_memories", + apiKey: "tpuf_xxxxxxxxxxxx", + region: "aws-us-west-2", + distanceMetric: "cosine_distance", + batchSize: 200, + }, + }, +}; +``` + diff --git a/mem0-ts/package.json b/mem0-ts/package.json index e8824e67e..717a23a07 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -119,6 +119,7 @@ "@pinecone-database/pinecone": "^8.0.0", "@qdrant/js-client-rest": "^1.18.0", "@supabase/supabase-js": "^2.49.1", + "@turbopuffer/turbopuffer": "^2.0.0", "@types/jest": "29.5.14", "@types/pg": "8.11.0", "better-sqlite3": "^12.6.2", diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 27a8d89bf..1da2a8096 100644 --- a/mem0-ts/pnpm-lock.yaml +++ b/mem0-ts/pnpm-lock.yaml @@ -65,6 +65,9 @@ importers: '@supabase/supabase-js': specifier: ^2.49.1 version: 2.108.1 + '@turbopuffer/turbopuffer': + specifier: ^2.0.0 + version: 2.5.0 '@types/jest': specifier: 29.5.14 version: 29.5.14 @@ -1284,6 +1287,9 @@ packages: '@tsconfig/node16@1.0.4': resolution: {integrity: sha512-vxhUy4J8lyeyinH7Azl1pdd43GJhZH/tP2weN8TntQblOY+A0XbT8DJk1/oCPuOOyg/Ja757rG0CgHcWC8OfMA==} + '@turbopuffer/turbopuffer@2.5.0': + resolution: {integrity: sha512-sUPvlynIaQNNtzeFgAttytZjGl430ttRMbXJq38vDPc+M7v/lo+K1M9/SpFimKMiyQsJXQiqSuhM85xdofr2fw==} + '@types/babel__core@7.20.5': resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==} @@ -2827,6 +2833,9 @@ packages: packet-reader@1.0.0: resolution: {integrity: sha512-HAKu/fG3HpHFO0AA8WE8q2g+gBJaZ9MG7fcKk+IJPLTGAD6Psw4443l+9DGRbOIh3/aXr7Phy0TjilYivJo5XQ==} + pako@2.2.0: + resolution: {integrity: sha512-zJq6RP/5q+TO2OpFV3FHzlPnFjmkb7Nc99a5SNjJE+uu/PkpChs+NIZSSzbBoD+6kjiISXjfYdwj1ZRQ81dz/w==} + parse-json@5.2.0: resolution: {integrity: sha512-ayCKvm/phCGxOkYRSCM82iDwct8/EonSEgCSxWxD7ve6jHggsFl4fZVQBPRNgQoKiuV/odhFrGzQXZwbifC8Rg==} engines: {node: '>=8'} @@ -5089,6 +5098,11 @@ snapshots: '@tsconfig/node16@1.0.4': {} + '@turbopuffer/turbopuffer@2.5.0': + dependencies: + pako: 2.2.0 + undici: 7.28.0 + '@types/babel__core@7.20.5': dependencies: '@babel/parser': 7.29.7 @@ -6866,6 +6880,8 @@ snapshots: packet-reader@1.0.0: {} + pako@2.2.0: {} + parse-json@5.2.0: dependencies: '@babel/code-frame': 7.29.7 diff --git a/mem0-ts/src/oss/src/tests/vector_stores/turbopuffer.test.ts b/mem0-ts/src/oss/src/tests/vector_stores/turbopuffer.test.ts new file mode 100644 index 000000000..b5cf40db4 --- /dev/null +++ b/mem0-ts/src/oss/src/tests/vector_stores/turbopuffer.test.ts @@ -0,0 +1,412 @@ +import { TurbopufferDB } from "../../vector_stores/turbopuffer"; + +// Mock the @turbopuffer/turbopuffer package +const mockWrite = jest.fn().mockResolvedValue(undefined); +const mockQuery = jest.fn(); +const mockDeleteAll = jest.fn().mockResolvedValue(undefined); + +const mockMigrationsNs = { + write: jest.fn().mockResolvedValue(undefined), + query: jest.fn(), + deleteAll: jest.fn().mockResolvedValue(undefined), +}; + +const mockNs = { + write: mockWrite, + query: mockQuery, + deleteAll: mockDeleteAll, +}; + +const mockNamespace = jest.fn((name: string) => { + if (name.endsWith("_migrations")) return mockMigrationsNs; + return mockNs; +}); + +jest.mock("@turbopuffer/turbopuffer", () => ({ + __esModule: true, + default: jest.fn().mockImplementation(() => ({ + namespace: mockNamespace, + })), +})); + +function makeStore(overrides: Record = {}): TurbopufferDB { + return new TurbopufferDB({ + apiKey: "test-key", + collectionName: "test-ns", + embeddingModelDims: 3, + ...overrides, + } as any); +} + +beforeEach(() => { + jest.clearAllMocks(); + mockWrite.mockResolvedValue(undefined); + mockQuery.mockResolvedValue({ rows: [] }); + mockDeleteAll.mockResolvedValue(undefined); + mockMigrationsNs.write.mockResolvedValue(undefined); + mockMigrationsNs.query.mockResolvedValue({ rows: [] }); + mockMigrationsNs.deleteAll.mockResolvedValue(undefined); +}); + +describe("TurbopufferDB constructor", () => { + it("throws when no API key is provided", () => { + const orig = process.env.TURBOPUFFER_API_KEY; + delete process.env.TURBOPUFFER_API_KEY; + expect( + () => + new TurbopufferDB({ + collectionName: "c", + embeddingModelDims: 3, + } as any), + ).toThrow(/API key/); + process.env.TURBOPUFFER_API_KEY = orig; + }); + + it("reads API key from environment variable", () => { + process.env.TURBOPUFFER_API_KEY = "env-key"; + expect( + () => + new TurbopufferDB({ + collectionName: "c", + embeddingModelDims: 3, + } as any), + ).not.toThrow(); + delete process.env.TURBOPUFFER_API_KEY; + }); +}); + +describe("TurbopufferDB initialize", () => { + it("resolves without calling write", async () => { + const store = makeStore(); + await expect(store.initialize()).resolves.toBeUndefined(); + expect(mockWrite).not.toHaveBeenCalled(); + }); +}); + +describe("TurbopufferDB keywordSearch", () => { + it("returns null", async () => { + const store = makeStore(); + expect(await store.keywordSearch()).toBeNull(); + }); +}); + +describe("TurbopufferDB insert", () => { + it("calls write with upsert_rows containing id and vector", async () => { + const store = makeStore(); + await store.insert([[0.1, 0.2, 0.3]], ["id-1"], [{ data: "hello" }]); + expect(mockWrite).toHaveBeenCalledTimes(1); + expect(mockWrite).toHaveBeenCalledWith({ + upsert_rows: [{ data: "hello", id: "id-1", vector: [0.1, 0.2, 0.3] }], + distance_metric: "cosine_distance", + }); + }); + + it("batches into multiple write calls when batchSize is exceeded", async () => { + const store = makeStore({ batchSize: 1 }); + await store.insert( + [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + ["id-1", "id-2"], + [{ a: 1 }, { b: 2 }], + ); + expect(mockWrite).toHaveBeenCalledTimes(2); + expect(mockWrite).toHaveBeenNthCalledWith(1, { + upsert_rows: [{ a: 1, id: "id-1", vector: [0.1, 0.2, 0.3] }], + distance_metric: "cosine_distance", + }); + expect(mockWrite).toHaveBeenNthCalledWith(2, { + upsert_rows: [{ b: 2, id: "id-2", vector: [0.4, 0.5, 0.6] }], + distance_metric: "cosine_distance", + }); + }); + + it("explicit id in payload is overridden by ids parameter", async () => { + const store = makeStore(); + await store.insert( + [[0.1, 0.2, 0.3]], + ["correct-id"], + [{ id: "wrong-id", data: "test" }], + ); + const call = mockWrite.mock.calls[0][0]; + expect(call.upsert_rows[0].id).toBe("correct-id"); + }); +}); + +describe("TurbopufferDB search", () => { + it("queries with correct rank_by, top_k, include_attributes", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0.1, 0.2, 0.3], 10); + expect(mockQuery).toHaveBeenCalledWith({ + rank_by: ["vector", "ANN", [0.1, 0.2, 0.3]], + top_k: 10, + include_attributes: true, + }); + }); + + it("defaults top_k to 5", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0.1, 0.2, 0.3]); + expect(mockQuery.mock.calls[0][0].top_k).toBe(5); + }); + + it("maps $dist to score as 1 - $dist and strips $dist and vector", async () => { + mockQuery.mockResolvedValue({ + rows: [{ id: "r1", $dist: 0.3, vector: [1, 2, 3], data: "payload-val" }], + }); + const store = makeStore(); + const results = await store.search([0.1, 0.2, 0.3], 5); + expect(results).toHaveLength(1); + expect(results[0].id).toBe("r1"); + expect(results[0].score).toBeCloseTo(0.7); + expect(results[0].payload).toEqual({ data: "payload-val" }); + expect(results[0].payload).not.toHaveProperty("$dist"); + expect(results[0].payload).not.toHaveProperty("vector"); + }); + + it("passes single-key filter as bare condition", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0.1, 0.2, 0.3], 5, { user_id: "u1" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters).toEqual(["user_id", "Eq", "u1"]); + }); + + it("passes multi-key filter as ['And', [...]]", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0.1, 0.2, 0.3], 5, { user_id: "u1", agent_id: "a1" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters[0]).toBe("And"); + expect(call.filters[1]).toContainEqual(["user_id", "Eq", "u1"]); + expect(call.filters[1]).toContainEqual(["agent_id", "Eq", "a1"]); + }); + + it("omits filters key when no filters provided", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0.1, 0.2, 0.3]); + const call = mockQuery.mock.calls[0][0]; + expect(call).not.toHaveProperty("filters"); + }); + + it("returns [] on query error", async () => { + mockQuery.mockRejectedValue(new Error("query failed")); + const store = makeStore(); + const results = await store.search([0.1, 0.2, 0.3]); + expect(results).toEqual([]); + }); +}); + +describe("TurbopufferDB get", () => { + it("queries with id filter and id-based rank, returns first row", async () => { + mockQuery.mockResolvedValue({ + rows: [{ id: "v1", $dist: 0.1, data: "stuff" }], + }); + const store = makeStore(); + const result = await store.get("v1"); + expect(result).not.toBeNull(); + expect(result!.id).toBe("v1"); + expect(result!.payload).toEqual({ data: "stuff" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters).toEqual(["id", "Eq", "v1"]); + expect(call.rank_by).toEqual(["id", "asc"]); + }); + + it("returns null when rows are empty", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + expect(await store.get("missing")).toBeNull(); + }); + + it("returns null on query error", async () => { + mockQuery.mockRejectedValue(new Error("network error")); + const store = makeStore(); + expect(await store.get("v1")).toBeNull(); + }); +}); + +describe("TurbopufferDB update", () => { + it("uses upsert_rows when vector is provided", async () => { + const store = makeStore(); + await store.update("id-1", [0.1, 0.2, 0.3], { data: "updated" }); + expect(mockWrite).toHaveBeenCalledWith({ + upsert_rows: [{ data: "updated", id: "id-1", vector: [0.1, 0.2, 0.3] }], + distance_metric: "cosine_distance", + }); + }); + + it("uses patch_rows when vector is empty", async () => { + const store = makeStore(); + await store.update("id-1", [], { data: "patched" }); + expect(mockWrite).toHaveBeenCalledWith({ + patch_rows: [{ data: "patched", id: "id-1" }], + }); + }); + + it("uses patch_rows when vector is null", async () => { + const store = makeStore(); + await store.update("id-1", null as any, { data: "patched" }); + expect(mockWrite).toHaveBeenCalledWith({ + patch_rows: [{ data: "patched", id: "id-1" }], + }); + }); +}); + +describe("TurbopufferDB delete", () => { + it("calls write with deletes array", async () => { + const store = makeStore(); + await store.delete("id-1"); + expect(mockWrite).toHaveBeenCalledWith({ deletes: ["id-1"] }); + }); +}); + +describe("TurbopufferDB deleteCol", () => { + it("calls deleteAll on the namespace", async () => { + const store = makeStore(); + await store.deleteCol(); + expect(mockDeleteAll).toHaveBeenCalledTimes(1); + }); +}); + +describe("TurbopufferDB list", () => { + it("lists rows ordered by id and returns [rows, count]", async () => { + mockQuery.mockResolvedValue({ + rows: [ + { id: "r1", $dist: 0.2, val: "a" }, + { id: "r2", $dist: 0.4, val: "b" }, + ], + }); + const store = makeStore(); + const [rows, count] = await store.list(); + expect(count).toBe(2); + expect(rows).toHaveLength(2); + expect(rows[0].id).toBe("r1"); + const call = mockQuery.mock.calls[0][0]; + expect(call.rank_by).toEqual(["id", "asc"]); + expect(call.top_k).toBe(100); + }); + + it("applies filters when provided", async () => { + mockQuery.mockResolvedValue({ rows: [{ id: "r1", $dist: 0.1 }] }); + const store = makeStore(); + await store.list({ user_id: "u1" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters).toEqual(["user_id", "Eq", "u1"]); + }); + + it("returns [[], 0] on error", async () => { + mockQuery.mockRejectedValue(new Error("list failed")); + const store = makeStore(); + const result = await store.list(); + expect(result).toEqual([[], 0]); + }); +}); + +describe("TurbopufferDB getUserId", () => { + it("generates a random ID and stores it when namespace is empty", async () => { + mockMigrationsNs.query.mockResolvedValue({ rows: [] }); + const store = makeStore(); + const id = await store.getUserId(); + expect(typeof id).toBe("string"); + expect(id.length).toBeGreaterThan(0); + expect(mockMigrationsNs.write).toHaveBeenCalledWith( + expect.objectContaining({ + upsert_rows: [expect.objectContaining({ id: "1", user_id: id })], + }), + ); + }); + + it("returns existing user_id without writing when found", async () => { + mockMigrationsNs.query.mockResolvedValue({ + rows: [{ id: "1", user_id: "existing-user" }], + }); + const store = makeStore(); + const id = await store.getUserId(); + expect(id).toBe("existing-user"); + expect(mockMigrationsNs.write).not.toHaveBeenCalled(); + }); + + it("creates a fresh id when the migrations namespace does not exist yet", async () => { + // A brand-new namespace 404s on the first read; getUserId should recover + // and persist a new id rather than surfacing the error. + mockMigrationsNs.query.mockRejectedValue( + Object.assign(new Error("namespace not found"), { status: 404 }), + ); + const store = makeStore(); + const id = await store.getUserId(); + expect(typeof id).toBe("string"); + expect(id.length).toBeGreaterThan(0); + expect(mockMigrationsNs.write).toHaveBeenCalledWith( + expect.objectContaining({ + upsert_rows: [expect.objectContaining({ id: "1", user_id: id })], + }), + ); + }); + + it("rethrows non-404 errors", async () => { + mockMigrationsNs.query.mockRejectedValue( + Object.assign(new Error("migrations error"), { status: 500 }), + ); + const store = makeStore(); + await expect(store.getUserId()).rejects.toThrow("migrations error"); + }); +}); + +describe("TurbopufferDB setUserId", () => { + it("writes upsert to migrations namespace", async () => { + const store = makeStore(); + await store.setUserId("my-user-id"); + expect(mockMigrationsNs.write).toHaveBeenCalledWith({ + upsert_rows: [{ id: "1", vector: [0.0], user_id: "my-user-id" }], + distance_metric: "cosine_distance", + }); + }); + + it("rethrows on error", async () => { + mockMigrationsNs.write.mockRejectedValue(new Error("write failed")); + const store = makeStore(); + await expect(store.setUserId("u")).rejects.toThrow("write failed"); + }); +}); + +describe("TurbopufferDB filter conversion", () => { + it("converts single Eq filter to bare condition", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0, 0, 0], 5, { status: "active" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters).toEqual(["status", "Eq", "active"]); + }); + + it("converts multi-field filters to ['And', [...]]", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0, 0, 0], 5, { user_id: "u1", type: "fact" }); + const call = mockQuery.mock.calls[0][0]; + expect(call.filters[0]).toBe("And"); + expect(call.filters[1]).toHaveLength(2); + }); + + it("converts range filter with gte/lte to multiple conditions", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0, 0, 0], 5, { score: { gte: 0.5, lte: 1.0 } } as any); + const call = mockQuery.mock.calls[0][0]; + // Two conditions wrapped in And + expect(call.filters[0]).toBe("And"); + expect(call.filters[1]).toContainEqual(["score", "Gte", 0.5]); + expect(call.filters[1]).toContainEqual(["score", "Lte", 1.0]); + }); + + it("omits filters key when filters object is empty", async () => { + mockQuery.mockResolvedValue({ rows: [] }); + const store = makeStore(); + await store.search([0, 0, 0], 5, {}); + const call = mockQuery.mock.calls[0][0]; + expect(call).not.toHaveProperty("filters"); + }); +}); diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 5a3e6b8e3..7caea89f3 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -46,6 +46,7 @@ import { VertexAIVectorSearch } from "../vector_stores/vertex_ai_vector_search"; import { CassandraDB } from "../vector_stores/cassandra"; import { PineconeDB } from "../vector_stores/pinecone"; import { S3Vectors } from "../vector_stores/s3_vectors"; +import { TurbopufferDB } from "../vector_stores/turbopuffer"; export class EmbedderFactory { static create(provider: string, config: EmbeddingConfig): Embedder { @@ -141,6 +142,8 @@ export class VectorStoreFactory { case "s3-vectors": case "s3_vectors": return new S3Vectors(config as any); + case "turbopuffer": + return new TurbopufferDB(config as any); default: throw new Error(`Unsupported vector store provider: ${provider}`); } diff --git a/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts b/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts new file mode 100644 index 000000000..e4bff3d7e --- /dev/null +++ b/mem0-ts/src/oss/src/vector_stores/turbopuffer.ts @@ -0,0 +1,238 @@ +import Turbopuffer from "@turbopuffer/turbopuffer"; +import { VectorStore } from "./base"; +import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types"; + +interface TurbopufferConfig extends VectorStoreConfig { + apiKey?: string; + region?: string; + collectionName: string; + distanceMetric?: string; + batchSize?: number; +} + +export class TurbopufferDB implements VectorStore { + private client: Turbopuffer; + private ns: ReturnType["namespace"]>; + private migrationsNs: ReturnType< + InstanceType["namespace"] + >; + private readonly collectionName: string; + private readonly distanceMetric: string; + private readonly batchSize: number; + + constructor(config: TurbopufferConfig) { + const apiKey = config.apiKey ?? process.env.TURBOPUFFER_API_KEY; + if (!apiKey) { + throw new Error( + "Turbopuffer API key is required. Provide it via config.apiKey or the TURBOPUFFER_API_KEY environment variable.", + ); + } + + this.client = new Turbopuffer({ + apiKey, + region: config.region ?? "gcp-us-central1", + }); + this.collectionName = config.collectionName; + this.distanceMetric = config.distanceMetric ?? "cosine_distance"; + this.batchSize = config.batchSize ?? 100; + this.ns = this.client.namespace(this.collectionName); + this.migrationsNs = this.client.namespace( + this.collectionName + "_migrations", + ); + } + + async initialize(): Promise { + // no-op: Turbopuffer creates namespaces on first write + } + + async insert( + vectors: number[][], + ids: string[], + payloads: Record[], + ): Promise { + for (let i = 0; i < vectors.length; i += this.batchSize) { + const batchVectors = vectors.slice(i, i + this.batchSize); + const batchIds = ids.slice(i, i + this.batchSize); + const batchPayloads = payloads.slice(i, i + this.batchSize); + + const upsert_rows = batchVectors.map((vector, j) => ({ + ...batchPayloads[j], + id: batchIds[j], + vector, + })); + + await this.ns.write({ + upsert_rows, + distance_metric: this.distanceMetric as any, + }); + } + } + + async search( + query: number[], + topK?: number, + filters?: SearchFilters, + ): Promise { + const queryParams: any = { + rank_by: ["vector", "ANN", query], + top_k: topK ?? 5, + include_attributes: true, + }; + + const tpufFilters = this.convertFilters(filters); + if (tpufFilters !== null) queryParams.filters = tpufFilters; + + try { + const result = await this.ns.query(queryParams); + return this.parseRows(result.rows ?? []); + } catch (err) { + console.error("Turbopuffer search error:", err); + return []; + } + } + + async keywordSearch(): Promise { + return null; + } + + async get(vectorId: string): Promise { + try { + const result = await this.ns.query({ + rank_by: ["id", "asc"] as any, + top_k: 1, + include_attributes: true, + filters: ["id", "Eq", vectorId] as any, + }); + const rows = result.rows ?? []; + return rows.length ? this.parseRows(rows)[0] : null; + } catch (err) { + console.error("Turbopuffer get error:", err); + return null; + } + } + + async update( + vectorId: string, + vector: number[], + payload: Record, + ): Promise { + if (vector && vector.length > 0) { + await this.ns.write({ + upsert_rows: [{ ...payload, id: vectorId, vector }], + distance_metric: this.distanceMetric as any, + }); + } else { + await this.ns.write({ + patch_rows: [{ ...payload, id: vectorId }], + }); + } + } + + async delete(vectorId: string): Promise { + await this.ns.write({ deletes: [vectorId] }); + } + + async deleteCol(): Promise { + await this.ns.deleteAll(); + } + + async list( + filters?: SearchFilters, + topK?: number, + ): Promise<[VectorStoreResult[], number]> { + const queryParams: any = { + rank_by: ["id", "asc"], + top_k: topK ?? 100, + include_attributes: true, + }; + + const tpufFilters = this.convertFilters(filters); + if (tpufFilters !== null) queryParams.filters = tpufFilters; + + try { + const result = await this.ns.query(queryParams); + const rows = this.parseRows(result.rows ?? []); + return [rows, rows.length]; + } catch (err) { + console.error("Turbopuffer list error:", err); + return [[], 0]; + } + } + + async getUserId(): Promise { + try { + let rows: any[] = []; + try { + const result = await this.migrationsNs.query({ + rank_by: ["id", "asc"] as any, + top_k: 1, + include_attributes: true, + }); + rows = result.rows ?? []; + } catch (err: any) { + // The migrations namespace is created lazily on first write, so the + // very first read 404s. Treat that as "no id yet" and fall through to + // create one; surface any other error (auth, rate limit, network). + if (err?.status !== 404) throw err; + } + if (rows.length > 0 && rows[0].user_id) { + return String(rows[0].user_id); + } + const randomId = + Math.random().toString(36).slice(2, 15) + + Math.random().toString(36).slice(2, 15); + await this.migrationsNs.write({ + upsert_rows: [{ id: "1", vector: [0.0], user_id: randomId }], + distance_metric: "cosine_distance" as any, + }); + return randomId; + } catch (err) { + console.error("Error getting user ID:", err); + throw err; + } + } + + async setUserId(userId: string): Promise { + try { + await this.migrationsNs.write({ + upsert_rows: [{ id: "1", vector: [0.0], user_id: userId }], + distance_metric: "cosine_distance" as any, + }); + } catch (err) { + console.error("Error setting user ID:", err); + throw err; + } + } + + private convertFilters(filters?: SearchFilters): any { + if (!filters || Object.keys(filters).length === 0) return null; + + const conditions: any[] = []; + for (const [key, value] of Object.entries(filters)) { + if ( + typeof value === "object" && + value !== null && + !Array.isArray(value) + ) { + if ("gte" in value) conditions.push([key, "Gte", value.gte]); + if ("lte" in value) conditions.push([key, "Lte", value.lte]); + if ("gt" in value) conditions.push([key, "Gt", value.gt]); + if ("lt" in value) conditions.push([key, "Lt", value.lt]); + } else { + conditions.push([key, "Eq", value]); + } + } + + if (conditions.length === 0) return null; + if (conditions.length === 1) return conditions[0]; + return ["And", conditions]; + } + + private parseRows(rows: any[]): VectorStoreResult[] { + return rows.map((row) => { + const { id, $dist, vector, ...rest } = row; + const score = $dist != null ? 1 - $dist : undefined; + return { id: String(id), payload: rest, score }; + }); + } +} diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index f4f7216ee..6cd4ed914 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -28,6 +28,7 @@ const external = [ "fastembed", "compromise", "natural", + "@turbopuffer/turbopuffer", ]; const define = {