feat(vector-stores): add Pinecone provider to TypeScript SDK (#5802)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
@@ -10,7 +10,8 @@ description: "Use Pinecone as a fully managed vector database in Mem0 with serve
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
import os
|
||||
from mem0 import Memory
|
||||
|
||||
@@ -44,10 +45,43 @@ messages = [
|
||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
import { Memory } from 'mem0ai/oss';
|
||||
|
||||
// Set OPENAI_API_KEY and PINECONE_API_KEY in your environment
|
||||
const config = {
|
||||
vectorStore: {
|
||||
provider: 'pinecone',
|
||||
config: {
|
||||
collectionName: 'testing',
|
||||
embeddingModelDims: 1536, // Matches OpenAI's text-embedding-3-small
|
||||
namespace: 'my-namespace', // Optional: specify a namespace for multi-tenancy
|
||||
serverlessConfig: {
|
||||
cloud: 'aws', // 'aws' | 'gcp' | 'azure'
|
||||
region: 'us-east-1',
|
||||
},
|
||||
metric: 'cosine',
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
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 thriller movies but I love sci-fi movies."},
|
||||
{"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
|
||||
]
|
||||
await memory.add(messages, { userId: "alice", metadata: { category: "movies" } });
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
### Config
|
||||
|
||||
Here are the parameters available for configuring Pinecone:
|
||||
|
||||
<Tabs>
|
||||
<Tab title="Python">
|
||||
| Parameter | Description | Default Value |
|
||||
| --- | --- | --- |
|
||||
| `collection_name` | Name of the index/collection | Required |
|
||||
@@ -61,11 +95,28 @@ Here are the parameters available for configuring Pinecone:
|
||||
| `metric` | Distance metric for vector similarity | `"cosine"` |
|
||||
| `batch_size` | Batch size for operations | `100` |
|
||||
| `namespace` | Namespace for the collection, useful for multi-tenancy. | `None` |
|
||||
</Tab>
|
||||
<Tab title="TypeScript">
|
||||
| Parameter | Description | Default Value |
|
||||
| --- | --- | --- |
|
||||
| `collectionName` | Name of the index/collection | Required |
|
||||
| `embeddingModelDims` | Dimensions of the embedding model (must match your chosen embedding model) | `1536` |
|
||||
| `client` | Existing Pinecone client instance | `undefined` |
|
||||
| `apiKey` | API key for Pinecone | Environment variable: `PINECONE_API_KEY` |
|
||||
| `serverlessConfig` | Configuration for serverless deployment (`cloud`, `region`) | `undefined` |
|
||||
| `podConfig` | Configuration for pod-based deployment (`environment`, `podType`, `pods`, `replicas`, `shards`) | `undefined` |
|
||||
| `metric` | Distance metric for vector similarity (`cosine`, `dotproduct`, `euclidean`) | `"cosine"` |
|
||||
| `batchSize` | Batch size for insert operations | `100` |
|
||||
| `namespace` | Namespace for the collection, useful for multi-tenancy. | `undefined` |
|
||||
| `extraParams` | Extra parameters spread into the Pinecone `createIndex` call | `{}` |
|
||||
</Tab>
|
||||
</Tabs>
|
||||
|
||||
> **Important**: You must choose either `serverless_config` or `pod_config` for your deployment, but not both.
|
||||
|
||||
#### Serverless Config Example
|
||||
```python
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "pinecone",
|
||||
@@ -82,8 +133,27 @@ config = {
|
||||
}
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
const config = {
|
||||
vectorStore: {
|
||||
provider: 'pinecone',
|
||||
config: {
|
||||
collectionName: 'memory_index',
|
||||
embeddingModelDims: 1536, // For OpenAI's text-embedding-3-small
|
||||
namespace: 'my-namespace', // Optional: custom namespace
|
||||
serverlessConfig: {
|
||||
cloud: 'aws', // 'gcp' | 'azure'
|
||||
region: 'us-east-1', // Choose appropriate region
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
#### Pod Config Example
|
||||
```python
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
config = {
|
||||
"vector_store": {
|
||||
"provider": "pinecone",
|
||||
@@ -99,4 +169,23 @@ config = {
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
const config = {
|
||||
vectorStore: {
|
||||
provider: 'pinecone',
|
||||
config: {
|
||||
collectionName: 'memory_index',
|
||||
embeddingModelDims: 1536, // For OpenAI's text-embedding-ada-002
|
||||
namespace: 'my-namespace', // Optional: custom namespace
|
||||
podConfig: {
|
||||
environment: 'gcp-starter',
|
||||
replicas: 1,
|
||||
podType: 'starter',
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
@@ -115,6 +115,7 @@
|
||||
"@google/genai": "^1.40.0",
|
||||
"@langchain/core": "^1.1.47",
|
||||
"@mistralai/mistralai": "^1.5.2",
|
||||
"@pinecone-database/pinecone": "^8.0.0",
|
||||
"@qdrant/js-client-rest": "^1.18.0",
|
||||
"@supabase/supabase-js": "^2.49.1",
|
||||
"@types/jest": "29.5.14",
|
||||
|
||||
Generated
+9
@@ -53,6 +53,9 @@ importers:
|
||||
'@mistralai/mistralai':
|
||||
specifier: ^1.5.2
|
||||
version: 1.15.1
|
||||
'@pinecone-database/pinecone':
|
||||
specifier: ^8.0.0
|
||||
version: 8.0.0
|
||||
'@qdrant/js-client-rest':
|
||||
specifier: ^1.18.0
|
||||
version: 1.18.0(typescript@5.5.4)
|
||||
@@ -796,6 +799,10 @@ packages:
|
||||
resolution: {integrity: sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==}
|
||||
engines: {node: '>= 8'}
|
||||
|
||||
'@pinecone-database/pinecone@8.0.0':
|
||||
resolution: {integrity: sha512-dItFqLdis2Pd5lC67aKn8HhvXajzlOz4+0RyK1CRcZMdSwG8YxUCBtH4yXPC8/6uLxfr2dvMWnRFuKgtPwwajQ==}
|
||||
engines: {node: '>=20.0.0'}
|
||||
|
||||
'@pkgjs/parseargs@0.11.0':
|
||||
resolution: {integrity: sha512-+1VkjdD0QBLPodGrJUeqarH8VAIvQODIbwh9XpP5Syisf7YoQgsJKPNFoqqLQlu+VQ/tVSshMR6loPMn8U+dPg==}
|
||||
engines: {node: '>=14'}
|
||||
@@ -4403,6 +4410,8 @@ snapshots:
|
||||
'@nodelib/fs.scandir': 2.1.5
|
||||
fastq: 1.20.1
|
||||
|
||||
'@pinecone-database/pinecone@8.0.0': {}
|
||||
|
||||
'@pkgjs/parseargs@0.11.0':
|
||||
optional: true
|
||||
|
||||
|
||||
@@ -0,0 +1,529 @@
|
||||
// jest.mock is hoisted before variable declarations, so we cannot close over
|
||||
// variables declared with let/const. All shared mock functions are attached to
|
||||
// the module-level `__mocks__` object that is populated inside the factory so
|
||||
// that the hoisted mock can reach them via a stable reference.
|
||||
|
||||
const __mocks__: {
|
||||
upsert: jest.Mock;
|
||||
query: jest.Mock;
|
||||
fetch: jest.Mock;
|
||||
deleteOne: jest.Mock;
|
||||
namespace: jest.Mock;
|
||||
describeIndexStats: jest.Mock;
|
||||
index: jest.Mock;
|
||||
listIndexes: jest.Mock;
|
||||
createIndex: jest.Mock;
|
||||
deleteIndex: jest.Mock;
|
||||
Pinecone: jest.Mock;
|
||||
} = {} as any;
|
||||
|
||||
jest.mock("@pinecone-database/pinecone", () => {
|
||||
// These are created fresh inside the factory so hoisting is safe.
|
||||
const upsert = jest.fn().mockResolvedValue(undefined);
|
||||
const query = jest.fn().mockResolvedValue({ matches: [] });
|
||||
const fetch = jest.fn().mockResolvedValue({ records: {} });
|
||||
const deleteOne = jest.fn().mockResolvedValue(undefined);
|
||||
|
||||
const nsHandle = { upsert, query, fetch, deleteOne };
|
||||
const namespace = jest.fn().mockReturnValue(nsHandle);
|
||||
|
||||
const describeIndexStats = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ totalRecordCount: 0, namespaces: {} });
|
||||
|
||||
const indexHandle = {
|
||||
namespace,
|
||||
describeIndexStats,
|
||||
// expose ops directly for the no-namespace path
|
||||
upsert,
|
||||
query,
|
||||
fetch,
|
||||
deleteOne,
|
||||
};
|
||||
const index = jest.fn().mockReturnValue(indexHandle);
|
||||
|
||||
const listIndexes = jest.fn().mockResolvedValue({ indexes: [] });
|
||||
const createIndex = jest.fn().mockResolvedValue(undefined);
|
||||
const deleteIndex = jest.fn().mockResolvedValue(undefined);
|
||||
|
||||
const Pinecone = jest.fn().mockImplementation(() => ({
|
||||
listIndexes,
|
||||
createIndex,
|
||||
deleteIndex,
|
||||
index,
|
||||
}));
|
||||
|
||||
// Populate the shared reference so tests can reach the mocks.
|
||||
Object.assign(__mocks__, {
|
||||
upsert,
|
||||
query,
|
||||
fetch,
|
||||
deleteOne,
|
||||
namespace,
|
||||
describeIndexStats,
|
||||
index,
|
||||
listIndexes,
|
||||
createIndex,
|
||||
deleteIndex,
|
||||
Pinecone,
|
||||
});
|
||||
|
||||
return { Pinecone };
|
||||
});
|
||||
|
||||
import { PineconeDB } from "../vector_stores/pinecone";
|
||||
import { VectorStoreFactory } from "../utils/factory";
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
function makeDb(overrides: Record<string, any> = {}): PineconeDB {
|
||||
return new PineconeDB({
|
||||
collectionName: "test-index",
|
||||
embeddingModelDims: 4,
|
||||
apiKey: "test-api-key",
|
||||
...overrides,
|
||||
} as any);
|
||||
}
|
||||
|
||||
async function initDb(
|
||||
overrides: Record<string, any> = {},
|
||||
): Promise<PineconeDB> {
|
||||
const db = makeDb(overrides);
|
||||
await db.initialize();
|
||||
return db;
|
||||
}
|
||||
|
||||
// --- Reset mocks between tests ---
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
|
||||
__mocks__.listIndexes.mockResolvedValue({ indexes: [] });
|
||||
__mocks__.createIndex.mockResolvedValue(undefined);
|
||||
__mocks__.deleteIndex.mockResolvedValue(undefined);
|
||||
__mocks__.upsert.mockResolvedValue(undefined);
|
||||
__mocks__.query.mockResolvedValue({ matches: [] });
|
||||
__mocks__.fetch.mockResolvedValue({ records: {} });
|
||||
__mocks__.deleteOne.mockResolvedValue(undefined);
|
||||
__mocks__.describeIndexStats.mockResolvedValue({
|
||||
totalRecordCount: 0,
|
||||
namespaces: {},
|
||||
});
|
||||
|
||||
const nsHandle = {
|
||||
upsert: __mocks__.upsert,
|
||||
query: __mocks__.query,
|
||||
fetch: __mocks__.fetch,
|
||||
deleteOne: __mocks__.deleteOne,
|
||||
};
|
||||
__mocks__.namespace.mockReturnValue(nsHandle);
|
||||
__mocks__.index.mockReturnValue({
|
||||
namespace: __mocks__.namespace,
|
||||
describeIndexStats: __mocks__.describeIndexStats,
|
||||
upsert: __mocks__.upsert,
|
||||
query: __mocks__.query,
|
||||
fetch: __mocks__.fetch,
|
||||
deleteOne: __mocks__.deleteOne,
|
||||
});
|
||||
__mocks__.Pinecone.mockImplementation(() => ({
|
||||
listIndexes: __mocks__.listIndexes,
|
||||
createIndex: __mocks__.createIndex,
|
||||
deleteIndex: __mocks__.deleteIndex,
|
||||
index: __mocks__.index,
|
||||
}));
|
||||
});
|
||||
|
||||
// --- Test suites ---
|
||||
|
||||
describe("VectorStoreFactory", () => {
|
||||
it("returns a PineconeDB instance for provider 'pinecone'", async () => {
|
||||
const db = VectorStoreFactory.create("pinecone", {
|
||||
collectionName: "x",
|
||||
embeddingModelDims: 4,
|
||||
apiKey: "k",
|
||||
} as any);
|
||||
expect(db).toBeInstanceOf(PineconeDB);
|
||||
await (db as any).initialize();
|
||||
});
|
||||
});
|
||||
|
||||
describe("Constructor", () => {
|
||||
it("uses apiKey from config", async () => {
|
||||
await initDb({ apiKey: "from-config" });
|
||||
expect(__mocks__.Pinecone).toHaveBeenCalledWith({ apiKey: "from-config" });
|
||||
});
|
||||
|
||||
it("falls back to PINECONE_API_KEY env var", async () => {
|
||||
process.env.PINECONE_API_KEY = "env-key";
|
||||
try {
|
||||
const db = new PineconeDB({
|
||||
collectionName: "test-index",
|
||||
embeddingModelDims: 4,
|
||||
} as any);
|
||||
await db.initialize();
|
||||
expect(__mocks__.Pinecone).toHaveBeenCalledWith({ apiKey: "env-key" });
|
||||
} finally {
|
||||
delete process.env.PINECONE_API_KEY;
|
||||
}
|
||||
});
|
||||
|
||||
it("throws when no API key is provided", () => {
|
||||
delete process.env.PINECONE_API_KEY;
|
||||
expect(
|
||||
() =>
|
||||
new PineconeDB({
|
||||
collectionName: "test-index",
|
||||
embeddingModelDims: 4,
|
||||
} as any),
|
||||
).toThrow("Pinecone API key required");
|
||||
});
|
||||
|
||||
it("accepts a pre-built client via config.client", async () => {
|
||||
const fakeClient = {
|
||||
listIndexes: __mocks__.listIndexes,
|
||||
createIndex: __mocks__.createIndex,
|
||||
deleteIndex: __mocks__.deleteIndex,
|
||||
index: __mocks__.index,
|
||||
};
|
||||
const db = new PineconeDB({
|
||||
collectionName: "test-index",
|
||||
embeddingModelDims: 4,
|
||||
client: fakeClient,
|
||||
} as any);
|
||||
await db.initialize();
|
||||
expect(__mocks__.Pinecone).not.toHaveBeenCalled();
|
||||
expect(__mocks__.listIndexes).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("initialize", () => {
|
||||
it("creates index with serverless default spec when index does not exist", async () => {
|
||||
await initDb();
|
||||
expect(__mocks__.createIndex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
name: "test-index",
|
||||
dimension: 4,
|
||||
metric: "cosine",
|
||||
spec: { serverless: { cloud: "aws", region: "us-east-1" } },
|
||||
waitUntilReady: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("creates index with pod spec when podConfig provided", async () => {
|
||||
await initDb({
|
||||
podConfig: { environment: "us-east1-gcp", podType: "p1.x2", pods: 2 },
|
||||
});
|
||||
expect(__mocks__.createIndex).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
spec: {
|
||||
pod: {
|
||||
environment: "us-east1-gcp",
|
||||
podType: "p1.x2",
|
||||
pods: 2,
|
||||
replicas: 1,
|
||||
shards: 1,
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("skips createIndex when index already exists", async () => {
|
||||
__mocks__.listIndexes.mockResolvedValue({
|
||||
indexes: [{ name: "test-index" }],
|
||||
});
|
||||
await initDb();
|
||||
expect(__mocks__.createIndex).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("_initPromise is shared across concurrent calls (idempotent)", async () => {
|
||||
// makeDb() fires initialize() in the constructor; calling it again before
|
||||
// it resolves must reuse the same in-flight promise so createIndex runs only once.
|
||||
const db = makeDb();
|
||||
await Promise.all([db.initialize(), db.initialize(), db.initialize()]);
|
||||
expect(__mocks__.createIndex).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("insert", () => {
|
||||
it("upserts records with correct shape", async () => {
|
||||
const db = await initDb();
|
||||
await db.insert([[1, 2, 3, 4]], ["id-1"], [{ text: "hello" }]);
|
||||
expect(__mocks__.upsert).toHaveBeenCalledWith({
|
||||
records: [
|
||||
{ id: "id-1", values: [1, 2, 3, 4], metadata: { text: "hello" } },
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
it("splits 150 records into two batches with batchSize=100", async () => {
|
||||
const db = await initDb({ batchSize: 100 });
|
||||
const vectors = Array.from({ length: 150 }, () => [0, 0, 0, 0]);
|
||||
const ids = Array.from({ length: 150 }, (_, i) => `id-${i}`);
|
||||
const payloads = Array.from({ length: 150 }, () => ({}));
|
||||
await db.insert(vectors, ids, payloads);
|
||||
expect(__mocks__.upsert).toHaveBeenCalledTimes(2);
|
||||
expect(__mocks__.upsert.mock.calls[0][0].records).toHaveLength(100);
|
||||
expect(__mocks__.upsert.mock.calls[1][0].records).toHaveLength(50);
|
||||
});
|
||||
});
|
||||
|
||||
describe("search", () => {
|
||||
it("calls query with correct args", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 10);
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
vector: [1, 2, 3, 4],
|
||||
topK: 10,
|
||||
includeMetadata: true,
|
||||
includeValues: false,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("translates equality filter to Pinecone $eq", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { user_id: "alice" });
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
filter: { user_id: { $eq: "alice" } },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("translates range filter to $gte/$lte", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { score: { gte: 0.5, lte: 1.0 } });
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
filter: { score: { $gte: 0.5, $lte: 1.0 } },
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("translates array filter to $in", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { tag: ["a", "b"] });
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ filter: { tag: { $in: ["a", "b"] } } }),
|
||||
);
|
||||
});
|
||||
|
||||
it("omits wildcard '*' from filter", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { user_id: "*" });
|
||||
const call = __mocks__.query.mock.calls[0][0];
|
||||
expect(call.filter).toBeUndefined();
|
||||
});
|
||||
|
||||
it("translates OR filter", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, {
|
||||
OR: [{ tag: "x" }, { tag: "y" }],
|
||||
});
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
filter: {
|
||||
$or: [{ tag: { $eq: "x" } }, { tag: { $eq: "y" } }],
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it("passes no filter when filters is empty", async () => {
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, {});
|
||||
const call = __mocks__.query.mock.calls[0][0];
|
||||
expect(call.filter).toBeUndefined();
|
||||
});
|
||||
|
||||
it("warns and skips NOT operator (unsupported by Pinecone)", async () => {
|
||||
const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {});
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { NOT: [{ tag: "x" }] });
|
||||
expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("NOT"));
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
|
||||
it("warns and skips contains operator", async () => {
|
||||
const warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {});
|
||||
const db = await initDb();
|
||||
await db.search([1, 2, 3, 4], 5, { tag: { contains: "foo" } });
|
||||
expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("contains"));
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
|
||||
it("throws on unsupported filter operator", async () => {
|
||||
const db = await initDb();
|
||||
await expect(
|
||||
db.search([1, 2, 3, 4], 5, { tag: { regex: "^foo" } } as any),
|
||||
).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("maps response matches to VectorStoreResult shape", async () => {
|
||||
__mocks__.query.mockResolvedValue({
|
||||
matches: [
|
||||
{ id: "v1", metadata: { text: "hi" }, score: 0.9 },
|
||||
{ id: "v2", metadata: { text: "bye" }, score: 0.7 },
|
||||
],
|
||||
});
|
||||
const db = await initDb();
|
||||
const results = await db.search([1, 2, 3, 4]);
|
||||
expect(results).toEqual([
|
||||
{ id: "v1", payload: { text: "hi" }, score: 0.9 },
|
||||
{ id: "v2", payload: { text: "bye" }, score: 0.7 },
|
||||
]);
|
||||
});
|
||||
|
||||
it("returns [] when matches is empty", async () => {
|
||||
__mocks__.query.mockResolvedValue({ matches: [] });
|
||||
const db = await initDb();
|
||||
const results = await db.search([1, 2, 3, 4]);
|
||||
expect(results).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("get", () => {
|
||||
it("returns VectorStoreResult when record found", async () => {
|
||||
__mocks__.fetch.mockResolvedValue({
|
||||
records: {
|
||||
"vec-1": { id: "vec-1", metadata: { text: "foo" } },
|
||||
},
|
||||
});
|
||||
const db = await initDb();
|
||||
const result = await db.get("vec-1");
|
||||
expect(result).toEqual({ id: "vec-1", payload: { text: "foo" } });
|
||||
});
|
||||
|
||||
it("returns null when record not found", async () => {
|
||||
__mocks__.fetch.mockResolvedValue({ records: {} });
|
||||
const db = await initDb();
|
||||
const result = await db.get("missing");
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("update", () => {
|
||||
it("upserts a single record", async () => {
|
||||
const db = await initDb();
|
||||
await db.update("vec-1", [1, 2, 3, 4], { text: "updated" });
|
||||
expect(__mocks__.upsert).toHaveBeenCalledWith({
|
||||
records: [
|
||||
{ id: "vec-1", values: [1, 2, 3, 4], metadata: { text: "updated" } },
|
||||
],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("delete", () => {
|
||||
it("calls deleteOne with the vectorId", async () => {
|
||||
const db = await initDb();
|
||||
await db.delete("vec-1");
|
||||
expect(__mocks__.deleteOne).toHaveBeenCalledWith({ id: "vec-1" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("deleteCol", () => {
|
||||
it("calls deleteIndex and resets internal state so re-init creates fresh index", async () => {
|
||||
const db = await initDb();
|
||||
await db.deleteCol();
|
||||
expect(__mocks__.deleteIndex).toHaveBeenCalledWith("test-index");
|
||||
// After deleteCol, _index and _initPromise reset; next initialize triggers createIndex again.
|
||||
__mocks__.listIndexes.mockResolvedValue({ indexes: [] });
|
||||
await db.initialize();
|
||||
expect(__mocks__.createIndex).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("list", () => {
|
||||
it("passes zero vector to query", async () => {
|
||||
const db = await initDb({ embeddingModelDims: 4 });
|
||||
await db.list();
|
||||
expect(__mocks__.query).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ vector: [0, 0, 0, 0] }),
|
||||
);
|
||||
});
|
||||
|
||||
it("returns the number of matches as the count", async () => {
|
||||
__mocks__.query.mockResolvedValue({
|
||||
matches: [
|
||||
{ id: "a", metadata: {}, score: 0 },
|
||||
{ id: "b", metadata: {}, score: 0 },
|
||||
],
|
||||
});
|
||||
const db = await initDb();
|
||||
const [results, count] = await db.list();
|
||||
expect(results).toHaveLength(2);
|
||||
expect(count).toBe(2);
|
||||
});
|
||||
|
||||
it("does not make an extra describeIndexStats round-trip", async () => {
|
||||
__mocks__.query.mockResolvedValue({ matches: [] });
|
||||
const db = await initDb();
|
||||
await db.list();
|
||||
expect(__mocks__.describeIndexStats).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("getUserId", () => {
|
||||
it("returns existing user_id from migrations namespace", async () => {
|
||||
__mocks__.fetch.mockResolvedValue({
|
||||
records: {
|
||||
"mem0-user-id": {
|
||||
id: "mem0-user-id",
|
||||
metadata: { user_id: "u-123" },
|
||||
},
|
||||
},
|
||||
});
|
||||
const db = await initDb();
|
||||
const uid = await db.getUserId();
|
||||
expect(uid).toBe("u-123");
|
||||
expect(__mocks__.namespace).toHaveBeenCalledWith("__mem0_migrations__");
|
||||
});
|
||||
|
||||
it("generates and upserts a new user_id when absent", async () => {
|
||||
__mocks__.fetch.mockResolvedValue({ records: {} });
|
||||
const db = await initDb();
|
||||
const uid = await db.getUserId();
|
||||
expect(typeof uid).toBe("string");
|
||||
expect(uid.length).toBeGreaterThan(0);
|
||||
expect(__mocks__.upsert).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
records: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
id: "mem0-user-id",
|
||||
metadata: { user_id: uid },
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
expect(__mocks__.namespace).toHaveBeenCalledWith("__mem0_migrations__");
|
||||
});
|
||||
});
|
||||
|
||||
describe("setUserId", () => {
|
||||
it("upserts with correct id, zero vector, and metadata", async () => {
|
||||
const db = await initDb({ embeddingModelDims: 4 });
|
||||
await db.setUserId("u-456");
|
||||
expect(__mocks__.upsert).toHaveBeenCalledWith({
|
||||
records: [
|
||||
{
|
||||
id: "mem0-user-id",
|
||||
values: [0, 0, 0, 0],
|
||||
metadata: { user_id: "u-456" },
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(__mocks__.namespace).toHaveBeenCalledWith("__mem0_migrations__");
|
||||
});
|
||||
});
|
||||
|
||||
describe("keywordSearch", () => {
|
||||
it("returns null", async () => {
|
||||
const db = await initDb();
|
||||
const result = await db.keywordSearch("hello", 5);
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -41,6 +41,7 @@ import { LangchainEmbedder } from "../embeddings/langchain";
|
||||
import { LangchainVectorStore } from "../vector_stores/langchain";
|
||||
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
||||
import { PGVector } from "../vector_stores/pgvector";
|
||||
import { PineconeDB } from "../vector_stores/pinecone";
|
||||
import { S3Vectors } from "../vector_stores/s3_vectors";
|
||||
|
||||
export class EmbedderFactory {
|
||||
@@ -126,6 +127,8 @@ export class VectorStoreFactory {
|
||||
return new AzureAISearch(config as any);
|
||||
case "pgvector":
|
||||
return new PGVector(config as any);
|
||||
case "pinecone":
|
||||
return new PineconeDB(config as any);
|
||||
case "s3-vectors":
|
||||
case "s3_vectors":
|
||||
return new S3Vectors(config as any);
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
import { Pinecone } from "@pinecone-database/pinecone";
|
||||
import type { Index } from "@pinecone-database/pinecone";
|
||||
import { VectorStore } from "./base";
|
||||
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
|
||||
|
||||
const MIGRATIONS_NAMESPACE = "__mem0_migrations__";
|
||||
const MIGRATIONS_RECORD_ID = "mem0-user-id";
|
||||
|
||||
interface PineconeDBConfig extends VectorStoreConfig {
|
||||
collectionName: string;
|
||||
embeddingModelDims: number;
|
||||
client?: Pinecone;
|
||||
apiKey?: string;
|
||||
serverlessConfig?: { cloud: string; region: string };
|
||||
podConfig?: {
|
||||
environment: string;
|
||||
podType?: string;
|
||||
pods?: number;
|
||||
replicas?: number;
|
||||
shards?: number;
|
||||
};
|
||||
metric?: "cosine" | "dotproduct" | "euclidean";
|
||||
batchSize?: number;
|
||||
namespace?: string;
|
||||
extraParams?: Record<string, any>;
|
||||
}
|
||||
|
||||
export class PineconeDB implements VectorStore {
|
||||
private client: Pinecone;
|
||||
private readonly collectionName: string;
|
||||
private readonly dimension: number;
|
||||
private readonly metric: "cosine" | "dotproduct" | "euclidean";
|
||||
private readonly batchSize: number;
|
||||
private readonly namespace: string;
|
||||
private readonly serverlessConfig?: { cloud: string; region: string };
|
||||
private readonly podConfig?: {
|
||||
environment: string;
|
||||
podType?: string;
|
||||
pods?: number;
|
||||
replicas?: number;
|
||||
shards?: number;
|
||||
};
|
||||
private readonly extraParams: Record<string, any>;
|
||||
private _index?: Index;
|
||||
private _initPromise?: Promise<void>;
|
||||
|
||||
constructor(config: PineconeDBConfig) {
|
||||
if (config.client) {
|
||||
this.client = config.client;
|
||||
} else {
|
||||
const apiKey = config.apiKey || process.env.PINECONE_API_KEY;
|
||||
if (!apiKey) {
|
||||
throw new Error(
|
||||
"Pinecone API key required: pass apiKey or set PINECONE_API_KEY env var",
|
||||
);
|
||||
}
|
||||
this.client = new Pinecone({ apiKey });
|
||||
}
|
||||
|
||||
this.collectionName = config.collectionName;
|
||||
this.dimension = config.embeddingModelDims || config.dimension || 1536;
|
||||
this.metric = config.metric || "cosine";
|
||||
this.batchSize = config.batchSize || 100;
|
||||
this.namespace = config.namespace || "";
|
||||
this.serverlessConfig = config.serverlessConfig;
|
||||
this.podConfig = config.podConfig;
|
||||
this.extraParams = config.extraParams || {};
|
||||
|
||||
this.initialize().catch(console.error);
|
||||
}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (!this._initPromise) {
|
||||
this._initPromise = this._doInitialize();
|
||||
}
|
||||
return this._initPromise;
|
||||
}
|
||||
|
||||
private async _doInitialize(): Promise<void> {
|
||||
await this._ensureIndex();
|
||||
this._index = this.client.index({ name: this.collectionName });
|
||||
}
|
||||
|
||||
private async _ensureIndex(): Promise<void> {
|
||||
const indexList = await this.client.listIndexes();
|
||||
const exists = ((indexList as any).indexes || []).some(
|
||||
(idx: { name: string }) => idx.name === this.collectionName,
|
||||
);
|
||||
if (exists) return;
|
||||
|
||||
const spec: Record<string, any> = this.podConfig
|
||||
? {
|
||||
pod: {
|
||||
environment: this.podConfig.environment,
|
||||
podType: this.podConfig.podType || "p1.x1",
|
||||
pods: this.podConfig.pods || 1,
|
||||
replicas: this.podConfig.replicas || 1,
|
||||
shards: this.podConfig.shards || 1,
|
||||
},
|
||||
}
|
||||
: {
|
||||
serverless: this.serverlessConfig || {
|
||||
cloud: "aws",
|
||||
region: "us-east-1",
|
||||
},
|
||||
};
|
||||
|
||||
await this.client.createIndex({
|
||||
name: this.collectionName,
|
||||
dimension: this.dimension,
|
||||
metric: this.metric,
|
||||
spec,
|
||||
waitUntilReady: true,
|
||||
...this.extraParams,
|
||||
});
|
||||
}
|
||||
|
||||
private index(): Index {
|
||||
return this._index!;
|
||||
}
|
||||
|
||||
private namespacedIndex(): Index {
|
||||
return this.namespace
|
||||
? this.index().namespace(this.namespace)
|
||||
: this.index();
|
||||
}
|
||||
|
||||
private migrationsIndex(): Index {
|
||||
return this.index().namespace(MIGRATIONS_NAMESPACE);
|
||||
}
|
||||
|
||||
private createFilter(
|
||||
filters?: SearchFilters,
|
||||
): Record<string, any> | undefined {
|
||||
if (!filters || Object.keys(filters).length === 0) return undefined;
|
||||
|
||||
const result: Record<string, any> = {};
|
||||
|
||||
for (const [key, value] of Object.entries(filters)) {
|
||||
if (value === undefined || value === null) continue;
|
||||
|
||||
if (key === "AND" || key === "$and") {
|
||||
result["$and"] = (value as SearchFilters[]).map(
|
||||
(sub) => this.createFilter(sub) || {},
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if (key === "OR" || key === "$or") {
|
||||
result["$or"] = (value as SearchFilters[]).map(
|
||||
(sub) => this.createFilter(sub) || {},
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if (key === "NOT" || key === "$not") {
|
||||
console.warn(
|
||||
"Filter operator 'NOT' is not supported by Pinecone metadata filters; skipping.",
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (value === "*") continue;
|
||||
|
||||
if (Array.isArray(value)) {
|
||||
result[key] = { $in: value };
|
||||
continue;
|
||||
}
|
||||
|
||||
if (typeof value === "object" && value !== null) {
|
||||
const pineconeOps: Record<string, any> = {};
|
||||
for (const [op, opVal] of Object.entries(value)) {
|
||||
switch (op) {
|
||||
case "eq":
|
||||
pineconeOps["$eq"] = opVal;
|
||||
break;
|
||||
case "ne":
|
||||
pineconeOps["$ne"] = opVal;
|
||||
break;
|
||||
case "gt":
|
||||
pineconeOps["$gt"] = opVal;
|
||||
break;
|
||||
case "gte":
|
||||
pineconeOps["$gte"] = opVal;
|
||||
break;
|
||||
case "lt":
|
||||
pineconeOps["$lt"] = opVal;
|
||||
break;
|
||||
case "lte":
|
||||
pineconeOps["$lte"] = opVal;
|
||||
break;
|
||||
case "in":
|
||||
pineconeOps["$in"] = opVal;
|
||||
break;
|
||||
case "nin":
|
||||
pineconeOps["$nin"] = opVal;
|
||||
break;
|
||||
case "contains":
|
||||
case "icontains":
|
||||
console.warn(
|
||||
`Filter operator '${op}' is not supported by Pinecone metadata filters; skipping.`,
|
||||
);
|
||||
break;
|
||||
default:
|
||||
throw new Error(
|
||||
`Unsupported filter operator '${op}' for Pinecone`,
|
||||
);
|
||||
}
|
||||
}
|
||||
if (Object.keys(pineconeOps).length > 0) {
|
||||
result[key] = pineconeOps;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
result[key] = { $eq: value };
|
||||
}
|
||||
|
||||
return Object.keys(result).length > 0 ? result : undefined;
|
||||
}
|
||||
|
||||
async insert(
|
||||
vectors: number[][],
|
||||
ids: string[],
|
||||
payloads: Record<string, any>[],
|
||||
): Promise<void> {
|
||||
await this.initialize();
|
||||
const records = vectors.map((values, i) => ({
|
||||
id: ids[i],
|
||||
values,
|
||||
metadata: payloads[i] || {},
|
||||
}));
|
||||
for (let i = 0; i < records.length; i += this.batchSize) {
|
||||
await this.namespacedIndex().upsert({
|
||||
records: records.slice(i, i + this.batchSize),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async keywordSearch(
|
||||
_query: string,
|
||||
_topK?: number,
|
||||
_filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[] | null> {
|
||||
return null;
|
||||
}
|
||||
|
||||
async search(
|
||||
query: number[],
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
await this.initialize();
|
||||
const filter = this.createFilter(filters);
|
||||
const response = await this.namespacedIndex().query({
|
||||
vector: query,
|
||||
topK,
|
||||
includeMetadata: true,
|
||||
includeValues: false,
|
||||
...(filter ? { filter } : {}),
|
||||
});
|
||||
return (response.matches || []).map((match: any) => ({
|
||||
id: match.id,
|
||||
payload: (match.metadata as Record<string, any>) || {},
|
||||
score: match.score,
|
||||
}));
|
||||
}
|
||||
|
||||
async get(vectorId: string): Promise<VectorStoreResult | null> {
|
||||
await this.initialize();
|
||||
const response = await this.namespacedIndex().fetch({ ids: [vectorId] });
|
||||
const record = (response.records || {})[vectorId];
|
||||
if (!record) return null;
|
||||
return {
|
||||
id: record.id,
|
||||
payload: (record.metadata as Record<string, any>) || {},
|
||||
};
|
||||
}
|
||||
|
||||
async update(
|
||||
vectorId: string,
|
||||
vector: number[],
|
||||
payload: Record<string, any>,
|
||||
): Promise<void> {
|
||||
await this.initialize();
|
||||
await this.namespacedIndex().upsert({
|
||||
records: [{ id: vectorId, values: vector, metadata: payload }],
|
||||
});
|
||||
}
|
||||
|
||||
async delete(vectorId: string): Promise<void> {
|
||||
await this.initialize();
|
||||
await this.namespacedIndex().deleteOne({ id: vectorId });
|
||||
}
|
||||
|
||||
async deleteCol(): Promise<void> {
|
||||
if (this._initPromise) {
|
||||
await this._initPromise.catch(() => {});
|
||||
}
|
||||
await this.client.deleteIndex(this.collectionName);
|
||||
this._index = undefined;
|
||||
this._initPromise = undefined;
|
||||
}
|
||||
|
||||
async list(
|
||||
filters?: SearchFilters,
|
||||
topK: number = 100,
|
||||
): Promise<[VectorStoreResult[], number]> {
|
||||
await this.initialize();
|
||||
const zeroVector = new Array(this.dimension).fill(0);
|
||||
const filter = this.createFilter(filters);
|
||||
const response = await this.namespacedIndex().query({
|
||||
vector: zeroVector,
|
||||
topK,
|
||||
includeMetadata: true,
|
||||
includeValues: false,
|
||||
...(filter ? { filter } : {}),
|
||||
});
|
||||
const results = (response.matches || []).map((match: any) => ({
|
||||
id: match.id,
|
||||
payload: (match.metadata as Record<string, any>) || {},
|
||||
score: match.score,
|
||||
}));
|
||||
return [results, results.length];
|
||||
}
|
||||
|
||||
async getUserId(): Promise<string> {
|
||||
await this.initialize();
|
||||
try {
|
||||
const response = await this.migrationsIndex().fetch({
|
||||
ids: [MIGRATIONS_RECORD_ID],
|
||||
});
|
||||
const record = (response.records || {})[MIGRATIONS_RECORD_ID];
|
||||
if (record?.metadata?.user_id) {
|
||||
return record.metadata.user_id as string;
|
||||
}
|
||||
} catch {
|
||||
// no record yet, fall through
|
||||
}
|
||||
const randomUserId =
|
||||
Math.random().toString(36).substring(2, 15) +
|
||||
Math.random().toString(36).substring(2, 15);
|
||||
await this.migrationsIndex().upsert({
|
||||
records: [
|
||||
{
|
||||
id: MIGRATIONS_RECORD_ID,
|
||||
values: new Array(this.dimension).fill(0),
|
||||
metadata: { user_id: randomUserId },
|
||||
},
|
||||
],
|
||||
});
|
||||
return randomUserId;
|
||||
}
|
||||
|
||||
async setUserId(userId: string): Promise<void> {
|
||||
await this.initialize();
|
||||
await this.migrationsIndex().upsert({
|
||||
records: [
|
||||
{
|
||||
id: MIGRATIONS_RECORD_ID,
|
||||
values: new Array(this.dimension).fill(0),
|
||||
metadata: { user_id: userId },
|
||||
},
|
||||
],
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ const external = [
|
||||
"pg",
|
||||
"zod",
|
||||
"better-sqlite3",
|
||||
"@pinecone-database/pinecone",
|
||||
"@qdrant/js-client-rest",
|
||||
"redis",
|
||||
"iovalkey",
|
||||
|
||||
Reference in New Issue
Block a user