diff --git a/docs/components/vectordbs/dbs/neon.mdx b/docs/components/vectordbs/dbs/neon.mdx index 20c91a569..257b1eb5d 100644 --- a/docs/components/vectordbs/dbs/neon.mdx +++ b/docs/components/vectordbs/dbs/neon.mdx @@ -21,49 +21,47 @@ from mem0 import Memory load_dotenv() config = { - "vector_store": { - "provider": "pgvector", - "config": { - "connection_string": os.environ["DATABASE_URL"], - "collection_name": "memories", - "embedding_model_dims": 1536, - "hnsw": True, - }, - }, +"vector_store": { +"provider": "pgvector", +"config": { +"connection_string": os.environ["DATABASE_URL"], +"collection_name": "memories", +"embedding_model_dims": 1536, +"hnsw": True, +}, +}, } m = Memory.from_config(config) 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."}, +{"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."}, ] m.add(messages, user_id="alice", metadata={"category": "movies"}) results = m.search( - "What movies should I recommend?", - filters={"user_id": "alice"}, +"What movies should I recommend?", +filters={"user_id": "alice"}, ) print(results) -``` + +```` ```typescript TypeScript import "dotenv/config"; import { Memory } from "mem0ai/oss"; -const databaseUrl = new URL(process.env.DATABASE_URL!); - const m = new Memory({ vectorStore: { provider: "pgvector", config: { - user: decodeURIComponent(databaseUrl.username), - password: decodeURIComponent(databaseUrl.password), - host: databaseUrl.hostname, - port: Number(databaseUrl.port || 5432), - dbname: databaseUrl.pathname.slice(1) || "neondb", + connectionString: process.env.DATABASE_URL!, + ssl: { + rejectUnauthorized: false, + }, collectionName: "memories", dimension: 1536, embeddingModelDims: 1536, @@ -89,7 +87,8 @@ const results = await m.search("What movies should I recommend?", { }); console.log(results); -``` +```` + ## SQL Migration @@ -116,20 +115,19 @@ DATABASE_URL=postgresql://user:password@ep-example.us-east-2.aws.neon.tech/neond | `sslmode` | PostgreSQL SSL mode. Use `require` for Neon. | Driver default | -The current Mem0 TypeScript `pgvector` adapter takes individual Postgres fields, -so parse `DATABASE_URL` before creating `Memory`. +Use the Neon `DATABASE_URL` directly with `connectionString`. Set `ssl` if your runtime needs an explicit TLS config object. + +| Parameter | Description | Default | +| -------------------- | ---------------------------------------------- | -------------- | +| `connectionString` | Neon Postgres connection string. | Required | +| `ssl` | Optional TLS settings passed directly to `pg`. | Driver default | +| `collectionName` | Name for the vector collection. | `memories` | +| `dimension` | Vector dimension for Mem0 config. | Auto-detected | +| `embeddingModelDims` | Embedding model dimensions for table creation. | Required | +| `hnsw` | Enables HNSW indexing. | `false` | + +**TLS note:** `ssl: true` is sufficient for most Neon connections since Neon uses valid certificates. Use `ssl: { rejectUnauthorized: false }` only when connecting through Neon's connection pooler on certain edge runtimes (e.g. Cloudflare Workers) that require it, or when your environment does not trust the Neon CA chain. -| Parameter | Description | Default | -| --- | --- | --- | -| `user` | Database user. | Required | -| `password` | Database password. | Required | -| `host` | Database host. | Required | -| `port` | Database port. | `5432` | -| `dbname` | Database name. | `vector_store` | -| `collectionName` | Name for the vector collection. | `memories` | -| `dimension` | Vector dimension for Mem0 config. | Auto-detected | -| `embeddingModelDims` | Embedding model dimensions for table creation. | Required | -| `hnsw` | Enables HNSW indexing. | `false` | diff --git a/docs/components/vectordbs/dbs/pgvector.mdx b/docs/components/vectordbs/dbs/pgvector.mdx index 9f59d7ebb..302f31200 100644 --- a/docs/components/vectordbs/dbs/pgvector.mdx +++ b/docs/components/vectordbs/dbs/pgvector.mdx @@ -2,6 +2,7 @@ title: "pgvector" description: "Use pgvector as a vector store in Mem0 for PostgreSQL-based vector similarity search with open-source simplicity." --- + [pgvector](https://github.com/pgvector/pgvector) is an open-source vector similarity search extension for Postgres. After connecting to Postgres, run `CREATE EXTENSION IF NOT EXISTS vector;` to create the vector extension. ### Usage @@ -14,41 +15,38 @@ from mem0 import Memory os.environ["OPENAI_API_KEY"] = "sk-xx" config = { - "vector_store": { - "provider": "pgvector", - "config": { - "user": "test", - "password": "123", - "host": "127.0.0.1", - "port": "5432", - } - } +"vector_store": { +"provider": "pgvector", +"config": { +"user": "test", +"password": "123", +"host": "127.0.0.1", +"port": "5432", +}, +} } m = Memory.from_config(config) 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."} +{"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."}, ] m.add(messages, user_id="alice", metadata={"category": "movies"}) -``` + +```` ```typescript TypeScript -import { Memory } from 'mem0ai/oss'; +import { Memory } from "mem0ai/oss"; const config = { vectorStore: { - provider: 'pgvector', + provider: "pgvector", config: { - collectionName: 'memories', + collectionName: "memories", embeddingModelDims: 1536, - user: 'test', - password: '123', - host: '127.0.0.1', - port: 5432, - dbname: 'vector_store', // Optional; TypeScript OSS defaults to `vector_store` when omitted + connectionString: "postgresql://test:123@localhost:5432/vector_store", diskann: false, // Optional, requires pgvectorscale extension hnsw: false, // Optional, for HNSW indexing }, @@ -57,37 +55,44 @@ const config = { 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."} -] + { 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 pgvector: -| Parameter | Description | Default Value | -| --- | --- | --- | -| `dbname` | The name of the database | `postgres` | -| `collection_name` | The name of the collection | `mem0` | -| `embedding_model_dims` | Dimensions of the embedding model | `1536` | -| `user` | User name to connect to the database | `None` | -| `password` | Password to connect to the database | `None` | -| `host` | The host where the Postgres server is running | `None` | -| `port` | The port where the Postgres server is running | `None` | -| `diskann` | Whether to use diskann for vector similarity search (requires pgvectorscale) | `True` | -| `hnsw` | Whether to use hnsw for vector similarity search | `False` | -| `sslmode` | SSL mode for PostgreSQL connection (e.g., 'require', 'prefer', 'disable') | `None` | -| `connection_string` | PostgreSQL connection string (overrides individual connection parameters) | `None` | -| `connection_pool` | psycopg2 connection pool object (overrides connection string and individual parameters) | `None` | +| Parameter | SDK | Description | Default Value | +| -------------------- | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------- | +| `connectionString` | TypeScript OSS | PostgreSQL connection string for direct connections. When set, Mem0 connects to the target database directly and skips the bootstrap `postgres` database flow. | `None` | +| `ssl` | TypeScript OSS | SSL option passed directly to `pg`, either `true` or an SSL config object, for both `connectionString` and split-field connections. | `None` | +| `dbname` | TypeScript OSS | Split-field database name. This is only used when `connectionString` is absent. | `vector_store` | +| `collectionName` | TypeScript OSS | Collection name. | `memories` | +| `embeddingModelDims` | TypeScript OSS | Dimensions of the embedding model. | Required | +| `user` | TypeScript OSS + Python | Database user for split-field connections. | `None` | +| `password` | TypeScript OSS + Python | Database password for split-field connections. | `None` | +| `host` | TypeScript OSS + Python | Database host for split-field connections. | `None` | +| `port` | TypeScript OSS + Python | Database port for split-field connections. | `None` | +| `diskann` | TypeScript OSS + Python | Whether to use DiskANN for vector similarity search, requires pgvectorscale. | `False` | +| `hnsw` | TypeScript OSS + Python | Whether to use HNSW for vector similarity search. | TypeScript OSS: `False`, Python: `True` | +| `connection_string` | Python only | PostgreSQL connection string, overrides individual connection parameters. | `None` | +| `sslmode` | Python only | SSL mode for PostgreSQL connections, such as `require`, `prefer`, or `disable`. | `None` | +| `connection_pool` | Python only | psycopg connection pool object, overrides connection string and individual connection parameters. | `None` | -**Note (TypeScript OSS):** If you omit `dbname`, the TypeScript client uses the database name `vector_store`. Python defaults to `postgres` for `dbname`, as in the table above. +**TypeScript OSS:** Use `connectionString` plus optional `ssl` for managed Postgres setups. If you omit `connectionString`, Mem0 falls back to split fields and uses `dbname`, `user`, `password`, `host`, `port`, and optional `ssl`. + +**Python:** The Python SDK uses snake_case keys such as `connection_string`, `sslmode`, `collection_name`, and `embedding_model_dims`. + +**Python connection priority**: -**Note**: The connection parameters have the following priority: 1. `connection_pool` (highest priority) 2. `connection_string` -3. Individual connection parameters (`user`, `password`, `host`, `port`, `sslmode`) \ No newline at end of file +3. Individual connection parameters (`user`, `password`, `host`, `port`, `sslmode`) diff --git a/mem0-ts/src/oss/src/vector_stores/pgvector.ts b/mem0-ts/src/oss/src/vector_stores/pgvector.ts index aba0060f6..450ffc9a5 100644 --- a/mem0-ts/src/oss/src/vector_stores/pgvector.ts +++ b/mem0-ts/src/oss/src/vector_stores/pgvector.ts @@ -1,4 +1,4 @@ -import type { Client as ClientType } from "pg"; +import type { Client as ClientType, ClientConfig } from "pg"; import pkg from "pg"; const { Client, escapeIdentifier } = pkg; import { VectorStore } from "./base"; @@ -157,41 +157,90 @@ export function buildFilterConditions( interface PGVectorConfig extends VectorStoreConfig { dbname?: string; - user: string; - password: string; - host: string; - port: number; + user?: string; + password?: string; + host?: string; + port?: number; + connectionString?: string; + ssl?: ClientConfig["ssl"]; embeddingModelDims: number; diskann?: boolean; hnsw?: boolean; } +function getConnectionString(config: PGVectorConfig): string | undefined { + return config.connectionString?.trim() || undefined; +} + +function validateConnectionConfig(config: PGVectorConfig): void { + if (getConnectionString(config)) { + return; + } + + const missingFields = ["user", "password", "host", "port"].filter((field) => { + const v = config[field as keyof PGVectorConfig]; + return v === undefined || v === null || v === ""; + }); + + if (missingFields.length > 0) { + throw new Error( + `PGVector requires either connectionString or ${missingFields.join(", ")}`, + ); + } +} + +function buildClientConfig( + config: PGVectorConfig, + database?: string, +): ClientConfig { + const connectionString = getConnectionString(config); + if (connectionString) { + return { + connectionString, + ...(config.ssl !== undefined ? { ssl: config.ssl } : {}), + }; + } + + return { + database, + user: config.user, + password: config.password, + host: config.host, + port: config.port, + ...(config.ssl !== undefined ? { ssl: config.ssl } : {}), + }; +} + export class PGVector implements VectorStore { private client: ClientType; private collectionName: string; private useDiskann: boolean; private useHnsw: boolean; private readonly dbName: string; + private readonly useDirectConnection: boolean; private config: PGVectorConfig; private _initPromise?: Promise; constructor(config: PGVectorConfig) { + validateConnectionConfig(config); this.collectionName = validateIdentifier( config.collectionName || "memories", "collectionName", ); this.useDiskann = config.diskann || false; this.useHnsw = config.hnsw || false; - this.dbName = validateIdentifier(config.dbname || "vector_store", "dbname"); + this.useDirectConnection = !!getConnectionString(config); + this.dbName = this.useDirectConnection + ? "" + : validateIdentifier(config.dbname || "vector_store", "dbname"); this.config = config; - this.client = new Client({ - database: "postgres", // Initially connect to default postgres database - user: config.user, - password: config.password, - host: config.host, - port: config.port, - }); + this.client = new Client( + buildClientConfig( + config, + this.useDirectConnection ? undefined : "postgres", + ), + ); this.initialize().catch(console.error); } @@ -210,29 +259,20 @@ export class PGVector implements VectorStore { try { await this.client.connect(); - // Check if database exists - const dbExists = await this.checkDatabaseExists(this.dbName); - if (!dbExists) { - await this.createDatabase(this.dbName); + if (!this.useDirectConnection) { + const dbExists = await this.checkDatabaseExists(this.dbName); + if (!dbExists) { + await this.createDatabase(this.dbName); + } + + await this.client.end(); + + this.client = new Client(buildClientConfig(this.config, this.dbName)); + await this.client.connect(); } - // Disconnect from postgres database - await this.client.end(); - - // Connect to the target database - this.client = new Client({ - database: this.dbName, - user: this.config.user, - password: this.config.password, - host: this.config.host, - port: this.config.port, - }); - await this.client.connect(); - - // Create vector extension await this.client.query("CREATE EXTENSION IF NOT EXISTS vector"); - // Create memory_migrations table await this.client.query(` CREATE TABLE IF NOT EXISTS memory_migrations ( id SERIAL PRIMARY KEY, @@ -240,7 +280,6 @@ export class PGVector implements VectorStore { ) `); - // Check if the collection exists const collections = await this.listCols(); if (!collections.includes(this.collectionName)) { await this.createCol(this.config.embeddingModelDims); diff --git a/mem0-ts/src/oss/tests/pgvector.filters.test.ts b/mem0-ts/src/oss/tests/pgvector.filters.test.ts index 6b6523be0..b46e33bb7 100644 --- a/mem0-ts/src/oss/tests/pgvector.filters.test.ts +++ b/mem0-ts/src/oss/tests/pgvector.filters.test.ts @@ -1,5 +1,3 @@ -/// - jest.mock("pg", () => { const Client = jest.fn().mockImplementation(() => ({ connect: jest.fn().mockResolvedValue(undefined), diff --git a/mem0-ts/src/oss/tests/pgvector.unit.test.ts b/mem0-ts/src/oss/tests/pgvector.unit.test.ts index d0d2b7cca..f6f75701f 100644 --- a/mem0-ts/src/oss/tests/pgvector.unit.test.ts +++ b/mem0-ts/src/oss/tests/pgvector.unit.test.ts @@ -1,5 +1,3 @@ -/// - const searchRows = [ { id: "a", @@ -23,9 +21,13 @@ const searchRows = [ }, ]; +const mockState = { + databaseExists: true, +}; + function mockPgQuery(sql: string) { if (sql.includes("SELECT 1 FROM pg_database")) { - return { rows: [{ "?column?": 1 }] }; + return { rows: mockState.databaseExists ? [{ "?column?": 1 }] : [] }; } if (sql.includes("FROM information_schema.tables")) { @@ -69,13 +71,129 @@ jest.mock("pg", () => { import { PGVector } from "../src/vector_stores/pgvector"; -describe("PGVector - search()", () => { +function getClientQueries(client: { query: jest.Mock }) { + return client.query.mock.calls.map(([sql]) => sql as string); +} + +describe("PGVector", () => { beforeEach(() => { const pg = require("pg"); + mockState.databaseExists = true; pg.__mock.Client.mockClear(); pg.__mock.clients.length = 0; }); + test("uses one direct client for connectionString mode and skips bootstrap database creation", async () => { + mockState.databaseExists = false; + + const ssl = { rejectUnauthorized: false }; + const store = new PGVector({ + collectionName: "memories", + connectionString: + "postgresql://postgres:postgres@db.example.com:5432/neondb", + ssl, + embeddingModelDims: 3, + dimension: 3, + } as any); + + await store.initialize(); + + const pg = require("pg"); + expect(pg.__mock.Client).toHaveBeenCalledTimes(1); + expect(pg.__mock.Client).toHaveBeenCalledWith({ + connectionString: + "postgresql://postgres:postgres@db.example.com:5432/neondb", + ssl, + }); + + const directClient = pg.__mock.clients[0]; + const queries = getClientQueries(directClient); + + expect(queries).not.toEqual( + expect.arrayContaining([ + expect.stringContaining("SELECT 1 FROM pg_database"), + ]), + ); + expect(queries).not.toEqual( + expect.arrayContaining([expect.stringContaining("CREATE DATABASE")]), + ); + expect(queries).toEqual( + expect.arrayContaining([ + "CREATE EXTENSION IF NOT EXISTS vector", + expect.stringContaining("FROM information_schema.tables"), + ]), + ); + }); + + test("keeps the split-field bootstrap flow when connectionString is absent", async () => { + mockState.databaseExists = false; + const ssl = { rejectUnauthorized: false }; + + const store = new PGVector({ + collectionName: "memories", + user: "postgres", + password: "postgres", + host: "localhost", + port: 5432, + dbname: "vector_store", + ssl, + embeddingModelDims: 3, + dimension: 3, + } as any); + + await store.initialize(); + + const pg = require("pg"); + expect(pg.__mock.Client).toHaveBeenCalledTimes(2); + expect(pg.__mock.Client).toHaveBeenNthCalledWith(1, { + database: "postgres", + user: "postgres", + password: "postgres", + host: "localhost", + port: 5432, + ssl, + }); + expect(pg.__mock.Client).toHaveBeenNthCalledWith(2, { + database: "vector_store", + user: "postgres", + password: "postgres", + host: "localhost", + port: 5432, + ssl, + }); + + const bootstrapClient = pg.__mock.clients[0]; + const activeClient = pg.__mock.clients[1]; + const bootstrapQueries = getClientQueries(bootstrapClient); + + expect(bootstrapQueries).toEqual( + expect.arrayContaining([ + "SELECT 1 FROM pg_database WHERE datname = $1", + 'CREATE DATABASE "vector_store"', + ]), + ); + expect(bootstrapClient.end).toHaveBeenCalledTimes(1); + expect(getClientQueries(activeClient)).toEqual( + expect.arrayContaining([ + "CREATE EXTENSION IF NOT EXISTS vector", + expect.stringContaining("FROM information_schema.tables"), + ]), + ); + }); + + test("throws when connectionString is absent and split-field params are missing", () => { + expect( + () => + new PGVector({ + collectionName: "memories", + embeddingModelDims: 3, + dimension: 3, + } as any), + ).toThrow( + "PGVector requires either connectionString or user, password, host, port", + ); + }); + test("returns similarity score (1 - distance) clamped to [0, 1]", async () => { const store = new PGVector({ collectionName: "memories",