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
|
### Usage
|
||||||
|
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
import os
|
import os
|
||||||
from mem0 import Memory
|
from mem0 import Memory
|
||||||
|
|
||||||
@@ -44,10 +45,43 @@ messages = [
|
|||||||
m.add(messages, user_id="alice", metadata={"category": "movies"})
|
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
|
### Config
|
||||||
|
|
||||||
Here are the parameters available for configuring Pinecone:
|
Here are the parameters available for configuring Pinecone:
|
||||||
|
|
||||||
|
<Tabs>
|
||||||
|
<Tab title="Python">
|
||||||
| Parameter | Description | Default Value |
|
| Parameter | Description | Default Value |
|
||||||
| --- | --- | --- |
|
| --- | --- | --- |
|
||||||
| `collection_name` | Name of the index/collection | Required |
|
| `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"` |
|
| `metric` | Distance metric for vector similarity | `"cosine"` |
|
||||||
| `batch_size` | Batch size for operations | `100` |
|
| `batch_size` | Batch size for operations | `100` |
|
||||||
| `namespace` | Namespace for the collection, useful for multi-tenancy. | `None` |
|
| `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.
|
> **Important**: You must choose either `serverless_config` or `pod_config` for your deployment, but not both.
|
||||||
|
|
||||||
#### Serverless Config Example
|
#### Serverless Config Example
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
config = {
|
config = {
|
||||||
"vector_store": {
|
"vector_store": {
|
||||||
"provider": "pinecone",
|
"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
|
#### Pod Config Example
|
||||||
```python
|
<CodeGroup>
|
||||||
|
```python Python
|
||||||
config = {
|
config = {
|
||||||
"vector_store": {
|
"vector_store": {
|
||||||
"provider": "pinecone",
|
"provider": "pinecone",
|
||||||
@@ -100,3 +170,22 @@ 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",
|
"@google/genai": "^1.40.0",
|
||||||
"@langchain/core": "^1.1.47",
|
"@langchain/core": "^1.1.47",
|
||||||
"@mistralai/mistralai": "^1.5.2",
|
"@mistralai/mistralai": "^1.5.2",
|
||||||
|
"@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",
|
||||||
"@types/jest": "29.5.14",
|
"@types/jest": "29.5.14",
|
||||||
|
|||||||
Generated
+9
@@ -53,6 +53,9 @@ importers:
|
|||||||
'@mistralai/mistralai':
|
'@mistralai/mistralai':
|
||||||
specifier: ^1.5.2
|
specifier: ^1.5.2
|
||||||
version: 1.15.1
|
version: 1.15.1
|
||||||
|
'@pinecone-database/pinecone':
|
||||||
|
specifier: ^8.0.0
|
||||||
|
version: 8.0.0
|
||||||
'@qdrant/js-client-rest':
|
'@qdrant/js-client-rest':
|
||||||
specifier: ^1.18.0
|
specifier: ^1.18.0
|
||||||
version: 1.18.0(typescript@5.5.4)
|
version: 1.18.0(typescript@5.5.4)
|
||||||
@@ -796,6 +799,10 @@ packages:
|
|||||||
resolution: {integrity: sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==}
|
resolution: {integrity: sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==}
|
||||||
engines: {node: '>= 8'}
|
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':
|
'@pkgjs/parseargs@0.11.0':
|
||||||
resolution: {integrity: sha512-+1VkjdD0QBLPodGrJUeqarH8VAIvQODIbwh9XpP5Syisf7YoQgsJKPNFoqqLQlu+VQ/tVSshMR6loPMn8U+dPg==}
|
resolution: {integrity: sha512-+1VkjdD0QBLPodGrJUeqarH8VAIvQODIbwh9XpP5Syisf7YoQgsJKPNFoqqLQlu+VQ/tVSshMR6loPMn8U+dPg==}
|
||||||
engines: {node: '>=14'}
|
engines: {node: '>=14'}
|
||||||
@@ -4403,6 +4410,8 @@ snapshots:
|
|||||||
'@nodelib/fs.scandir': 2.1.5
|
'@nodelib/fs.scandir': 2.1.5
|
||||||
fastq: 1.20.1
|
fastq: 1.20.1
|
||||||
|
|
||||||
|
'@pinecone-database/pinecone@8.0.0': {}
|
||||||
|
|
||||||
'@pkgjs/parseargs@0.11.0':
|
'@pkgjs/parseargs@0.11.0':
|
||||||
optional: true
|
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 { LangchainVectorStore } from "../vector_stores/langchain";
|
||||||
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
import { AzureAISearch } from "../vector_stores/azure_ai_search";
|
||||||
import { PGVector } from "../vector_stores/pgvector";
|
import { PGVector } from "../vector_stores/pgvector";
|
||||||
|
import { PineconeDB } from "../vector_stores/pinecone";
|
||||||
import { S3Vectors } from "../vector_stores/s3_vectors";
|
import { S3Vectors } from "../vector_stores/s3_vectors";
|
||||||
|
|
||||||
export class EmbedderFactory {
|
export class EmbedderFactory {
|
||||||
@@ -126,6 +127,8 @@ export class VectorStoreFactory {
|
|||||||
return new AzureAISearch(config as any);
|
return new AzureAISearch(config as any);
|
||||||
case "pgvector":
|
case "pgvector":
|
||||||
return new PGVector(config as any);
|
return new PGVector(config as any);
|
||||||
|
case "pinecone":
|
||||||
|
return new PineconeDB(config as any);
|
||||||
case "s3-vectors":
|
case "s3-vectors":
|
||||||
case "s3_vectors":
|
case "s3_vectors":
|
||||||
return new S3Vectors(config as any);
|
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",
|
"pg",
|
||||||
"zod",
|
"zod",
|
||||||
"better-sqlite3",
|
"better-sqlite3",
|
||||||
|
"@pinecone-database/pinecone",
|
||||||
"@qdrant/js-client-rest",
|
"@qdrant/js-client-rest",
|
||||||
"redis",
|
"redis",
|
||||||
"iovalkey",
|
"iovalkey",
|
||||||
|
|||||||
Reference in New Issue
Block a user