@mastra/pg 0.1.4 → 0.1.5-alpha.1

package/.turbo/turbo-build.log

@@ -1,18 +1,18 @@
 
-> @mastra/pg@0.1.4 build /home/runner/work/mastra/mastra/stores/pg
+> @mastra/pg@0.1.5-alpha.1 build /home/runner/work/mastra/mastra/stores/pg
 > tsup src/index.ts --format esm --experimental-dts --clean --treeshake
 
 CLI Building entry: src/index.ts
 CLI Using tsconfig: tsconfig.json
 CLI tsup v8.3.6
 TSC Build start
-TSC ⚡️ Build success in 8503ms
+TSC ⚡️ Build success in 9745ms
 DTS Build start
 CLI Target: es2022
 Analysis will use the bundled TypeScript version 5.7.3
 Writing package typings: /home/runner/work/mastra/mastra/stores/pg/dist/_tsup-dts-rollup.d.ts
-DTS ⚡️ Build success in 6464ms
+DTS ⚡️ Build success in 6324ms
 CLI Cleaning output folder
 ESM Build start
-ESM dist/index.js 30.26 KB
-ESM ⚡️ Build success in 584ms
+ESM dist/index.js 32.67 KB
+ESM ⚡️ Build success in 1005ms

package/CHANGELOG.md

@@ -1,5 +1,21 @@
 # @mastra/pg
 
+## 0.1.5-alpha.1
+
+### Patch Changes
+
+- 82197f8: Update PG vector to allow for multiple index types
+- Updated dependencies [6cb63e0]
+  - @mastra/core@0.4.2-alpha.1
+
+## 0.1.5-alpha.0
+
+### Patch Changes
+
+- Updated dependencies [7fceae1]
+- Updated dependencies [f626fbb]
+  - @mastra/core@0.4.2-alpha.0
+
 ## 0.1.4
 
 ### Patch Changes

package/README.md

@@ -56,7 +56,7 @@ const store = new PostgresStore({
   port: 5432,
   database: 'mastra',
   user: 'postgres',
-  password: 'postgres'
+  password: 'postgres',
 });
 
 // Create a thread
@@ -64,17 +64,19 @@ await store.saveThread({
   id: 'thread-123',
   resourceId: 'resource-456',
   title: 'My Thread',
-  metadata: { key: 'value' }
+  metadata: { key: 'value' },
 });
 
 // Add messages to thread
-await store.saveMessages([{
-  id: 'msg-789',
-  threadId: 'thread-123',
-  role: 'user',
-  type: 'text',
-  content: [{ type: 'text', text: 'Hello' }]
-}]);
+await store.saveMessages([
+  {
+    id: 'msg-789',
+    threadId: 'thread-123',
+    role: 'user',
+    type: 'text',
+    content: [{ type: 'text', text: 'Hello' }],
+  },
+]);
 
 // Query threads and messages
 const savedThread = await store.getThread('thread-123');
@@ -97,14 +99,18 @@ Connection pool settings:
 ## Features
 
 ### Vector Store Features
+
 - Vector similarity search with cosine, euclidean, and dot product metrics
 - Advanced metadata filtering with MongoDB-like query syntax
 - Minimum score threshold for queries
 - Automatic UUID generation for vectors
 - Table management (create, list, describe, delete, truncate)
-- Uses pgvector's IVFFLAT indexing with 100 lists
+- Uses pgvector's IVFFLAT indexing with 100 lists by default
+- Supports HNSW indexing with configurable parameters
+- Supports flat indexing
 
 ### Storage Features
+
 - Thread and message storage with JSON support
 - Atomic transactions for data consistency
 - Efficient batch operations

package/dist/_tsup-dts-rollup.d.ts

@@ -19,8 +19,47 @@ import { StorageThreadType } from '@mastra/core/memory';
 import { TABLE_NAMES } from '@mastra/core/storage';
 import { WorkflowRunState } from '@mastra/core/workflows';
 
+export declare const baseTestConfigs: {
+    smokeTests: {
+        dimension: number;
+        size: number;
+        k: number;
+        queryCount: number;
+    }[];
+    '64': {
+        dimension: number;
+        size: number;
+        k: number;
+        queryCount: number;
+    }[];
+    '384': {
+        dimension: number;
+        size: number;
+        k: number;
+        queryCount: number;
+    }[];
+    '1024': {
+        dimension: number;
+        size: number;
+        k: number;
+        queryCount: number;
+    }[];
+    stressTests: {
+        dimension: number;
+        size: number;
+        k: number;
+        queryCount: number;
+    }[];
+};
+
 export declare function buildFilterQuery(filter: Filter, minScore: number): FilterResult;
 
+export declare const calculateRecall: (actual: number[], expected: number[], k: number) => number;
+
+export declare const calculateTimeout: (dimension: number, size: number, k: number) => number;
+
+export declare function cosineSimilarity(a: number[], b: number[]): number;
+
 export declare const FILTER_OPERATORS: Record<string, OperatorFn>;
 
 declare type FilterOperator = {
@@ -34,8 +73,60 @@ export declare interface FilterResult {
     values: any[];
 }
 
+export declare const findNearestBruteForce: (query: number[], vectors: number[][], k: number) => number[];
+
+export declare const formatTable: (data: any[], columns: string[]) => string;
+
+export declare const generateClusteredVectors: (count: number, dim: number, numClusters?: number) => number[][];
+
+export declare const generateRandomVectors: (count: number, dim: number) => number[][];
+
+export declare const generateSkewedVectors: (count: number, dim: number) => number[][];
+
+export declare const getHNSWConfig: (indexConfig: IndexConfig) => {
+    m: number;
+    efConstruction: number;
+};
+
+export declare function getIndexDescription({ type, hnsw, }: {
+    type: IndexType;
+    hnsw: {
+        m: number;
+        efConstruction: number;
+    };
+}): string;
+
+export declare const getListCount: (indexConfig: IndexConfig, size: number) => number | undefined;
+
+export declare function getSearchEf(k: number, m: number): {
+    default: number;
+    lower: number;
+    higher: number;
+};
+
+export declare const groupBy: <T, K extends keyof T>(array: T[], key: K | ((item: T) => string), reducer?: (group: T[]) => any) => Record<string, any>;
+
 export declare const handleKey: (key: string) => string;
 
+declare interface HNSWConfig {
+    m?: number;
+    efConstruction?: number;
+}
+
+export declare interface IndexConfig {
+    type?: IndexType;
+    ivf?: IVFConfig;
+    hnsw?: HNSWConfig;
+}
+
+export declare type IndexType = 'ivfflat' | 'hnsw' | 'flat';
+
+declare interface IVFConfig {
+    lists?: number;
+}
+
+export declare function measureLatency<T>(fn: () => Promise<T>): Promise<[number, T]>;
+
 declare type OperatorFn = (key: string, paramIndex: number, value?: any) => FilterOperator;
 
 export declare type OperatorType = BasicOperator | NumericOperator | ArrayOperator | ElementOperator | LogicalOperator | '$contains' | Exclude<RegexOperator, '$options'>;
@@ -57,15 +148,34 @@ export declare class PGFilterTranslator extends BaseFilterTranslator {
     private translateRegexPattern;
 }
 
+declare interface PGIndexStats extends IndexStats {
+    type: IndexType;
+    config: {
+        m?: number;
+        efConstruction?: number;
+        lists?: number;
+        probes?: number;
+    };
+}
+export { PGIndexStats }
+export { PGIndexStats as PGIndexStats_alias_1 }
+
 declare class PgVector extends MastraVector {
     private pool;
+    private indexCache;
     constructor(connectionString: string);
     transformFilter(filter?: Filter): Filter;
-    query(indexName: string, queryVector: number[], topK?: number, filter?: Filter, includeVector?: boolean, minScore?: number): Promise<QueryResult[]>;
+    getIndexInfo(indexName: string): Promise<PGIndexStats>;
+    query(indexName: string, queryVector: number[], topK?: number, filter?: Filter, includeVector?: boolean, minScore?: number, // Optional minimum score threshold
+    options?: {
+        ef?: number;
+        probes?: number;
+    }): Promise<QueryResult[]>;
     upsert(indexName: string, vectors: number[][], metadata?: Record<string, any>[], ids?: string[]): Promise<string[]>;
-    createIndex(indexName: string, dimension: number, metric?: 'cosine' | 'euclidean' | 'dotproduct'): Promise<void>;
+    createIndex(indexName: string, dimension: number, metric?: 'cosine' | 'euclidean' | 'dotproduct', indexConfig?: IndexConfig, defineIndex?: boolean): Promise<void>;
+    defineIndex(indexName: string, metric: "cosine" | "euclidean" | "dotproduct" | undefined, indexConfig: IndexConfig): Promise<void>;
     listIndexes(): Promise<string[]>;
-    describeIndex(indexName: string): Promise<IndexStats>;
+    describeIndex(indexName: string): Promise<PGIndexStats>;
     deleteIndex(indexName: string): Promise<void>;
     truncateIndex(indexName: string): Promise<void>;
     disconnect(): Promise<void>;
@@ -151,4 +261,40 @@ declare class PostgresStore extends MastraStorage {
 export { PostgresStore }
 export { PostgresStore as PostgresStore_alias_1 }
 
+export declare interface TestConfig {
+    dimension: number;
+    size: number;
+    k: number;
+    queryCount: number;
+}
+
+export declare interface TestResult {
+    distribution: string;
+    dimension: number;
+    type: IndexType;
+    size: number;
+    k?: number;
+    metrics: {
+        recall?: number;
+        minRecall?: number;
+        maxRecall?: number;
+        latency?: {
+            p50: number;
+            p95: number;
+            lists?: number;
+            vectorsPerList?: number;
+            m?: number;
+            ef?: number;
+        };
+        clustering?: {
+            numLists?: number;
+            avgVectorsPerList?: number;
+            recommendedLists?: number;
+            distribution?: string;
+        };
+    };
+}
+
+export declare function warmupQuery(vectorDB: PgVector, indexName: string, dimension: number, k: number): Promise<void>;
+
 export { }

package/dist/index.d.ts

@@ -1,3 +1,4 @@
+export { PGIndexStats } from './_tsup-dts-rollup.js';
 export { PgVector } from './_tsup-dts-rollup.js';
 export { PostgresConfig } from './_tsup-dts-rollup.js';
 export { PostgresStore } from './_tsup-dts-rollup.js';

package/dist/index.js

@@ -286,6 +286,7 @@ function buildFilterQuery(filter, minScore) {
 // src/vector/index.ts
 var PgVector = class extends MastraVector {
   pool;
+  indexCache = /* @__PURE__ */ new Map();
   constructor(connectionString) {
     super();
     const basePool = new pg.Pool({
@@ -310,12 +311,27 @@ var PgVector = class extends MastraVector {
     const translatedFilter = pgFilter.translate(filter ?? {});
     return translatedFilter;
   }
-  async query(indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0) {
+  async getIndexInfo(indexName) {
+    if (!this.indexCache.has(indexName)) {
+      this.indexCache.set(indexName, await this.describeIndex(indexName));
+    }
+    return this.indexCache.get(indexName);
+  }
+  async query(indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0, options) {
     const client = await this.pool.connect();
     try {
       const vectorStr = `[${queryVector.join(",")}]`;
       const translatedFilter = this.transformFilter(filter);
       const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
+      const indexInfo = await this.getIndexInfo(indexName);
+      if (indexInfo.type === "hnsw") {
+        const calculatedEf = options?.ef ?? Math.max(topK, (indexInfo?.config?.m ?? 16) * topK);
+        const searchEf = Math.min(1e3, Math.max(1, calculatedEf));
+        await client.query(`SET LOCAL hnsw.ef_search = ${searchEf}`);
+      }
+      if (indexInfo.type === "ivfflat" && options?.probes) {
+        await client.query(`SET LOCAL ivfflat.probes = ${options.probes}`);
+      }
       const query = `
         WITH vector_scores AS (
           SELECT
@@ -349,13 +365,13 @@ var PgVector = class extends MastraVector {
       const vectorIds = ids || vectors.map(() => crypto.randomUUID());
       for (let i = 0; i < vectors.length; i++) {
         const query = `
-            INSERT INTO ${indexName} (vector_id, embedding, metadata)
-            VALUES ($1, $2::vector, $3::jsonb)
-            ON CONFLICT (vector_id)
-            DO UPDATE SET
-                embedding = $2::vector,
-                metadata = $3::jsonb
-            RETURNING embedding::text
+          INSERT INTO ${indexName} (vector_id, embedding, metadata)
+          VALUES ($1, $2::vector, $3::jsonb)
+          ON CONFLICT (vector_id)
+          DO UPDATE SET
+            embedding = $2::vector,
+            metadata = $3::jsonb
+          RETURNING embedding::text
         `;
         await client.query(query, [vectorIds[i], `[${vectors[i]?.join(",")}]`, JSON.stringify(metadata?.[i] || {})]);
       }
@@ -368,7 +384,7 @@ var PgVector = class extends MastraVector {
       client.release();
     }
   }
-  async createIndex(indexName, dimension, metric = "cosine") {
+  async createIndex(indexName, dimension, metric = "cosine", indexConfig = {}, defineIndex = true) {
     const client = await this.pool.connect();
     try {
       if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
@@ -394,13 +410,9 @@ var PgVector = class extends MastraVector {
           metadata JSONB DEFAULT '{}'::jsonb
         );
       `);
-      const indexMethod = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
-      await client.query(`
-        CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx
-        ON public.${indexName}
-        USING ivfflat (embedding ${indexMethod})
-        WITH (lists = 100);
-      `);
+      if (defineIndex) {
+        await this.defineIndex(indexName, metric, indexConfig);
+      }
     } catch (error) {
       console.error("Failed to create vector table:", error);
       throw error;
@@ -408,6 +420,46 @@ var PgVector = class extends MastraVector {
       client.release();
     }
   }
+  async defineIndex(indexName, metric = "cosine", indexConfig) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(`DROP INDEX IF EXISTS ${indexName}_vector_idx`);
+      if (indexConfig.type === "flat") return;
+      const metricOp = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
+      let indexSQL;
+      if (indexConfig.type === "hnsw") {
+        const m = indexConfig.hnsw?.m ?? 8;
+        const efConstruction = indexConfig.hnsw?.efConstruction ?? 32;
+        indexSQL = `
+          CREATE INDEX ${indexName}_vector_idx 
+          ON ${indexName} 
+          USING hnsw (embedding ${metricOp})
+          WITH (
+            m = ${m},
+            ef_construction = ${efConstruction}
+          )
+        `;
+      } else {
+        let lists;
+        if (indexConfig.ivf?.lists) {
+          lists = indexConfig.ivf.lists;
+        } else {
+          const size = (await client.query(`SELECT COUNT(*) FROM ${indexName}`)).rows[0].count;
+          lists = Math.max(100, Math.min(4e3, Math.floor(Math.sqrt(size) * 2)));
+        }
+        indexSQL = `
+          CREATE INDEX ${indexName}_vector_idx
+          ON ${indexName}
+          USING ivfflat (embedding ${metricOp})
+          WITH (lists = ${lists});
+        `;
+      }
+      await client.query(indexSQL);
+      this.indexCache.delete(indexName);
+    } finally {
+      client.release();
+    }
+  }
   async listIndexes() {
     const client = await this.pool.connect();
     try {
@@ -436,9 +488,10 @@ var PgVector = class extends MastraVector {
                 SELECT COUNT(*) as count
                 FROM ${indexName};
             `;
-      const metricQuery = `
+      const indexQuery = `
             SELECT
                 am.amname as index_method,
+                pg_get_indexdef(i.indexrelid) as index_def,
                 opclass.opcname as operator_class
             FROM pg_index i
             JOIN pg_class c ON i.indexrelid = c.oid
@@ -446,26 +499,33 @@ var PgVector = class extends MastraVector {
             JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
             WHERE c.relname = '${indexName}_vector_idx';
             `;
-      const [dimResult, countResult, metricResult] = await Promise.all([
+      const [dimResult, countResult, indexResult] = await Promise.all([
         client.query(dimensionQuery, [indexName]),
         client.query(countQuery),
-        client.query(metricQuery)
+        client.query(indexQuery)
       ]);
-      let metric = "cosine";
-      if (metricResult.rows.length > 0) {
-        const operatorClass = metricResult.rows[0].operator_class;
-        if (operatorClass.includes("l2")) {
-          metric = "euclidean";
-        } else if (operatorClass.includes("ip")) {
-          metric = "dotproduct";
-        } else if (operatorClass.includes("cosine")) {
-          metric = "cosine";
-        }
+      const { index_method, index_def, operator_class } = indexResult.rows[0] || {
+        index_method: "flat",
+        index_def: "",
+        operator_class: "cosine"
+      };
+      const metric = operator_class.includes("l2") ? "euclidean" : operator_class.includes("ip") ? "dotproduct" : "cosine";
+      const config = {};
+      if (index_method === "hnsw") {
+        const m = index_def.match(/m\s*=\s*'?(\d+)'?/)?.[1];
+        const efConstruction = index_def.match(/ef_construction\s*=\s*'?(\d+)'?/)?.[1];
+        if (m) config.m = parseInt(m);
+        if (efConstruction) config.efConstruction = parseInt(efConstruction);
+      } else if (index_method === "ivfflat") {
+        const lists = index_def.match(/lists\s*=\s*'?(\d+)'?/)?.[1];
+        if (lists) config.lists = parseInt(lists);
       }
       return {
         dimension: dimResult.rows[0].dimension,
         count: parseInt(countResult.rows[0].count),
-        metric
+        metric,
+        type: index_method,
+        config
       };
     } catch (e) {
       await client.query("ROLLBACK");

package/docker-compose.perf.yaml

@@ -0,0 +1,21 @@
+services:
+  db:
+    image: pgvector/pgvector:pg16
+    container_name: 'pg-perf-test-db'
+    ports:
+      - '5435:5432' 
+    environment:
+      POSTGRES_USER: ${POSTGRES_USER:-postgres}
+      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-postgres}
+      POSTGRES_DB: ${POSTGRES_DB:-mastra}
+    shm_size: 1gb
+    command:
+      - "postgres"
+      - "-c"
+      - "shared_buffers=512MB"
+      - "-c"
+      - "maintenance_work_mem=1024MB"
+      - "-c"
+      - "work_mem=512MB"
+    tmpfs:
+      - /var/lib/postgresql/data 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/pg",
-  "version": "0.1.4",
+  "version": "0.1.5-alpha.1",
   "description": "Postgres provider for Mastra - includes both vector and db storage capabilities",
   "type": "module",
   "main": "dist/index.js",
@@ -17,7 +17,7 @@
   "dependencies": {
     "pg": "^8.13.1",
     "pg-promise": "^11.5.4",
-    "@mastra/core": "^0.4.1"
+    "@mastra/core": "^0.4.2-alpha.1"
   },
   "devDependencies": {
     "@microsoft/api-extractor": "^7.49.2",
@@ -32,7 +32,10 @@
     "build:watch": "pnpm build --watch",
     "pretest": "docker compose up -d && (for i in $(seq 1 30); do docker compose exec -T db pg_isready -U postgres && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
     "test": "vitest run",
+    "pretest:perf": "docker compose -f docker-compose.perf.yaml up -d && (for i in $(seq 1 30); do docker compose -f docker-compose.perf.yaml exec -T db pg_isready -U postgres && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
+    "test:perf": "NODE_OPTIONS='--max-old-space-size=16384' vitest run -c vitest.perf.config.ts",
     "posttest": "docker compose down -v",
+    "posttest:perf": "docker compose -f docker-compose.perf.yaml down -v",
     "pretest:watch": "docker compose up -d",
     "test:watch": "vitest watch",
     "posttest:watch": "docker compose down -v"