feat(vector-stores): add Turbopuffer provider to TypeScript OSS SDK (#5801)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

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<string, any> = {}): 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");
+  });
+});

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}`);
     }

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<InstanceType<typeof Turbopuffer>["namespace"]>;
+  private migrationsNs: ReturnType<
+    InstanceType<typeof Turbopuffer>["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<void> {
+    // no-op: Turbopuffer creates namespaces on first write
+  }
+
+  async insert(
+    vectors: number[][],
+    ids: string[],
+    payloads: Record<string, any>[],
+  ): Promise<void> {
+    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<VectorStoreResult[]> {
+    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<null> {
+    return null;
+  }
+
+  async get(vectorId: string): Promise<VectorStoreResult | null> {
+    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<string, any>,
+  ): Promise<void> {
+    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<void> {
+    await this.ns.write({ deletes: [vectorId] });
+  }
+
+  async deleteCol(): Promise<void> {
+    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<string> {
+    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<void> {
+    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 };
+    });
+  }
+}