diff --git a/docs/components/vectordbs/dbs/pinecone.mdx b/docs/components/vectordbs/dbs/pinecone.mdx index 7911225fc..31662ed29 100644 --- a/docs/components/vectordbs/dbs/pinecone.mdx +++ b/docs/components/vectordbs/dbs/pinecone.mdx @@ -10,7 +10,8 @@ description: "Use Pinecone as a fully managed vector database in Mem0 with serve ### Usage -```python + +```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" } }); +``` + + ### Config Here are the parameters available for configuring Pinecone: + + | 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` | + + +| 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 | `{}` | + + > **Important**: You must choose either `serverless_config` or `pod_config` for your deployment, but not both. #### Serverless Config Example -```python + +```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 + }, + }, + }, +}; +``` + + #### Pod Config Example -```python + +```python Python config = { "vector_store": { "provider": "pinecone", @@ -99,4 +169,23 @@ config = { } } } -``` \ No newline at end of file +``` + +```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', + }, + }, + }, +}; +``` + diff --git a/mem0-ts/package.json b/mem0-ts/package.json index 84a471be4..2a25304e8 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -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", diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 399900522..8b9c59b50 100644 --- a/mem0-ts/pnpm-lock.yaml +++ b/mem0-ts/pnpm-lock.yaml @@ -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 diff --git a/mem0-ts/src/oss/src/tests/pinecone.test.ts b/mem0-ts/src/oss/src/tests/pinecone.test.ts new file mode 100644 index 000000000..0a266978d --- /dev/null +++ b/mem0-ts/src/oss/src/tests/pinecone.test.ts @@ -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 = {}): PineconeDB { + return new PineconeDB({ + collectionName: "test-index", + embeddingModelDims: 4, + apiKey: "test-api-key", + ...overrides, + } as any); +} + +async function initDb( + overrides: Record = {}, +): Promise { + 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(); + }); +}); diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 9ca9c26ae..076436527 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -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); diff --git a/mem0-ts/src/oss/src/vector_stores/pinecone.ts b/mem0-ts/src/oss/src/vector_stores/pinecone.ts new file mode 100644 index 000000000..4e8ef7212 --- /dev/null +++ b/mem0-ts/src/oss/src/vector_stores/pinecone.ts @@ -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; +} + +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; + private _index?: Index; + private _initPromise?: Promise; + + 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 { + if (!this._initPromise) { + this._initPromise = this._doInitialize(); + } + return this._initPromise; + } + + private async _doInitialize(): Promise { + await this._ensureIndex(); + this._index = this.client.index({ name: this.collectionName }); + } + + private async _ensureIndex(): Promise { + 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 = 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 | undefined { + if (!filters || Object.keys(filters).length === 0) return undefined; + + const result: Record = {}; + + 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 = {}; + 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[], + ): Promise { + 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 { + return null; + } + + async search( + query: number[], + topK: number = 5, + filters?: SearchFilters, + ): Promise { + 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) || {}, + score: match.score, + })); + } + + async get(vectorId: string): Promise { + 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) || {}, + }; + } + + async update( + vectorId: string, + vector: number[], + payload: Record, + ): Promise { + await this.initialize(); + await this.namespacedIndex().upsert({ + records: [{ id: vectorId, values: vector, metadata: payload }], + }); + } + + async delete(vectorId: string): Promise { + await this.initialize(); + await this.namespacedIndex().deleteOne({ id: vectorId }); + } + + async deleteCol(): Promise { + 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) || {}, + score: match.score, + })); + return [results, results.length]; + } + + async getUserId(): Promise { + 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 { + await this.initialize(); + await this.migrationsIndex().upsert({ + records: [ + { + id: MIGRATIONS_RECORD_ID, + values: new Array(this.dimension).fill(0), + metadata: { user_id: userId }, + }, + ], + }); + } +} diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index 23c9ab7dc..47ced0207 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -10,6 +10,7 @@ const external = [ "pg", "zod", "better-sqlite3", + "@pinecone-database/pinecone", "@qdrant/js-client-rest", "redis", "iovalkey",