feat(vector-stores): add Turbopuffer provider to TypeScript OSS SDK (#5801)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
@@ -6,7 +6,8 @@ description: "Use Turbopuffer as a serverless vector database in Mem0 for low-la
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
<CodeGroup>
|
||||
```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" });
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
### 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` |
|
||||
|
||||
<Note>
|
||||
**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.
|
||||
</Note>
|
||||
|
||||
### Regions
|
||||
|
||||
| Region | Location |
|
||||
@@ -62,7 +97,8 @@ Here are the parameters available for configuring Turbopuffer:
|
||||
|
||||
### Config Example
|
||||
|
||||
```python
|
||||
<CodeGroup>
|
||||
```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,
|
||||
},
|
||||
},
|
||||
};
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
@@ -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",
|
||||
|
||||
Generated
+16
@@ -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
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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}`);
|
||||
}
|
||||
|
||||
@@ -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 };
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,7 @@ const external = [
|
||||
"fastembed",
|
||||
"compromise",
|
||||
"natural",
|
||||
"@turbopuffer/turbopuffer",
|
||||
];
|
||||
|
||||
const define = {
|
||||
|
||||
Reference in New Issue
Block a user