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
|
### Usage
|
||||||
|
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
import os
|
import os
|
||||||
from mem0 import Memory
|
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"})
|
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
|
### Config
|
||||||
|
|
||||||
Here are the parameters available for configuring Turbopuffer:
|
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` |
|
| `batch_size` | Batch size for bulk operations | `100` |
|
||||||
| `extra_params` | Additional parameters for the Turbopuffer client | `None` |
|
| `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
|
### Regions
|
||||||
|
|
||||||
| Region | Location |
|
| Region | Location |
|
||||||
@@ -62,7 +97,8 @@ Here are the parameters available for configuring Turbopuffer:
|
|||||||
|
|
||||||
### Config Example
|
### Config Example
|
||||||
|
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
config = {
|
config = {
|
||||||
"vector_store": {
|
"vector_store": {
|
||||||
"provider": "turbopuffer",
|
"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",
|
"@pinecone-database/pinecone": "^8.0.0",
|
||||||
"@qdrant/js-client-rest": "^1.18.0",
|
"@qdrant/js-client-rest": "^1.18.0",
|
||||||
"@supabase/supabase-js": "^2.49.1",
|
"@supabase/supabase-js": "^2.49.1",
|
||||||
|
"@turbopuffer/turbopuffer": "^2.0.0",
|
||||||
"@types/jest": "29.5.14",
|
"@types/jest": "29.5.14",
|
||||||
"@types/pg": "8.11.0",
|
"@types/pg": "8.11.0",
|
||||||
"better-sqlite3": "^12.6.2",
|
"better-sqlite3": "^12.6.2",
|
||||||
|
|||||||
Generated
+16
@@ -65,6 +65,9 @@ importers:
|
|||||||
'@supabase/supabase-js':
|
'@supabase/supabase-js':
|
||||||
specifier: ^2.49.1
|
specifier: ^2.49.1
|
||||||
version: 2.108.1
|
version: 2.108.1
|
||||||
|
'@turbopuffer/turbopuffer':
|
||||||
|
specifier: ^2.0.0
|
||||||
|
version: 2.5.0
|
||||||
'@types/jest':
|
'@types/jest':
|
||||||
specifier: 29.5.14
|
specifier: 29.5.14
|
||||||
version: 29.5.14
|
version: 29.5.14
|
||||||
@@ -1284,6 +1287,9 @@ packages:
|
|||||||
'@tsconfig/node16@1.0.4':
|
'@tsconfig/node16@1.0.4':
|
||||||
resolution: {integrity: sha512-vxhUy4J8lyeyinH7Azl1pdd43GJhZH/tP2weN8TntQblOY+A0XbT8DJk1/oCPuOOyg/Ja757rG0CgHcWC8OfMA==}
|
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':
|
'@types/babel__core@7.20.5':
|
||||||
resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==}
|
resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==}
|
||||||
|
|
||||||
@@ -2827,6 +2833,9 @@ packages:
|
|||||||
packet-reader@1.0.0:
|
packet-reader@1.0.0:
|
||||||
resolution: {integrity: sha512-HAKu/fG3HpHFO0AA8WE8q2g+gBJaZ9MG7fcKk+IJPLTGAD6Psw4443l+9DGRbOIh3/aXr7Phy0TjilYivJo5XQ==}
|
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:
|
parse-json@5.2.0:
|
||||||
resolution: {integrity: sha512-ayCKvm/phCGxOkYRSCM82iDwct8/EonSEgCSxWxD7ve6jHggsFl4fZVQBPRNgQoKiuV/odhFrGzQXZwbifC8Rg==}
|
resolution: {integrity: sha512-ayCKvm/phCGxOkYRSCM82iDwct8/EonSEgCSxWxD7ve6jHggsFl4fZVQBPRNgQoKiuV/odhFrGzQXZwbifC8Rg==}
|
||||||
engines: {node: '>=8'}
|
engines: {node: '>=8'}
|
||||||
@@ -5089,6 +5098,11 @@ snapshots:
|
|||||||
|
|
||||||
'@tsconfig/node16@1.0.4': {}
|
'@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':
|
'@types/babel__core@7.20.5':
|
||||||
dependencies:
|
dependencies:
|
||||||
'@babel/parser': 7.29.7
|
'@babel/parser': 7.29.7
|
||||||
@@ -6866,6 +6880,8 @@ snapshots:
|
|||||||
|
|
||||||
packet-reader@1.0.0: {}
|
packet-reader@1.0.0: {}
|
||||||
|
|
||||||
|
pako@2.2.0: {}
|
||||||
|
|
||||||
parse-json@5.2.0:
|
parse-json@5.2.0:
|
||||||
dependencies:
|
dependencies:
|
||||||
'@babel/code-frame': 7.29.7
|
'@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 { CassandraDB } from "../vector_stores/cassandra";
|
||||||
import { PineconeDB } from "../vector_stores/pinecone";
|
import { PineconeDB } from "../vector_stores/pinecone";
|
||||||
import { S3Vectors } from "../vector_stores/s3_vectors";
|
import { S3Vectors } from "../vector_stores/s3_vectors";
|
||||||
|
import { TurbopufferDB } from "../vector_stores/turbopuffer";
|
||||||
|
|
||||||
export class EmbedderFactory {
|
export class EmbedderFactory {
|
||||||
static create(provider: string, config: EmbeddingConfig): Embedder {
|
static create(provider: string, config: EmbeddingConfig): Embedder {
|
||||||
@@ -141,6 +142,8 @@ export class VectorStoreFactory {
|
|||||||
case "s3-vectors":
|
case "s3-vectors":
|
||||||
case "s3_vectors":
|
case "s3_vectors":
|
||||||
return new S3Vectors(config as any);
|
return new S3Vectors(config as any);
|
||||||
|
case "turbopuffer":
|
||||||
|
return new TurbopufferDB(config as any);
|
||||||
default:
|
default:
|
||||||
throw new Error(`Unsupported vector store provider: ${provider}`);
|
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",
|
"fastembed",
|
||||||
"compromise",
|
"compromise",
|
||||||
"natural",
|
"natural",
|
||||||
|
"@turbopuffer/turbopuffer",
|
||||||
];
|
];
|
||||||
|
|
||||||
const define = {
|
const define = {
|
||||||
|
|||||||
Reference in New Issue
Block a user