@mastra/pg 0.1.6-alpha.0 → 0.1.6-alpha.3

package/src/vector/index.ts

@@ -1,10 +1,18 @@
-import type { Filter } from '@mastra/core/filter';
-import { type IndexStats, type QueryResult, MastraVector } from '@mastra/core/vector';
+import { MastraVector } from '@mastra/core/vector';
+import type {
+  IndexStats,
+  QueryResult,
+  QueryVectorParams,
+  CreateIndexParams,
+  UpsertVectorParams,
+  ParamsToArgs,
+} from '@mastra/core/vector';
+import type { VectorFilter } from '@mastra/core/vector/filter';
 import pg from 'pg';
 
 import { PGFilterTranslator } from './filter';
 import { buildFilterQuery } from './sql-builder';
-import { type IndexConfig, type IndexType } from './types';
+import type { IndexConfig, IndexType } from './types';
 
 export interface PGIndexStats extends IndexStats {
   type: IndexType;
@@ -16,6 +24,31 @@ export interface PGIndexStats extends IndexStats {
   };
 }
 
+interface PgQueryVectorParams extends QueryVectorParams {
+  minScore?: number;
+  /**
+   * HNSW search parameter. Controls the size of the dynamic candidate
+   * list during search. Higher values improve accuracy at the cost of speed.
+   */
+  ef?: number;
+  /**
+   * IVFFlat probe parameter. Number of cells to visit during search.
+   * Higher values improve accuracy at the cost of speed.
+   */
+  probes?: number;
+}
+
+interface PgCreateIndexParams extends CreateIndexParams {
+  indexConfig?: IndexConfig;
+  buildIndex?: boolean;
+}
+
+interface PgDefineIndexParams {
+  indexName: string;
+  metric: 'cosine' | 'euclidean' | 'dotproduct';
+  indexConfig: IndexConfig;
+}
+
 export class PgVector extends MastraVector {
   private pool: pg.Pool;
   private indexCache: Map<string, PGIndexStats> = new Map();
@@ -41,10 +74,9 @@ export class PgVector extends MastraVector {
       }) ?? basePool;
   }
 
-  transformFilter(filter?: Filter) {
-    const pgFilter = new PGFilterTranslator();
-    const translatedFilter = pgFilter.translate(filter ?? {});
-    return translatedFilter;
+  transformFilter(filter?: VectorFilter) {
+    const translator = new PGFilterTranslator();
+    return translator.translate(filter);
   }
 
   async getIndexInfo(indexName: string): Promise<PGIndexStats> {
@@ -54,18 +86,10 @@ export class PgVector extends MastraVector {
     return this.indexCache.get(indexName)!;
   }
 
-  async query(
-    indexName: string,
-    queryVector: number[],
-    topK: number = 10,
-    filter?: Filter,
-    includeVector: boolean = false,
-    minScore: number = 0, // Optional minimum score threshold
-    options?: {
-      ef?: number; // For HNSW
-      probes?: number; // For IVF
-    },
-  ): Promise<QueryResult[]> {
+  async query(...args: ParamsToArgs<PgQueryVectorParams>): Promise<QueryResult[]> {
+    const params = this.normalizeArgs<PgQueryVectorParams>('query', args, ['minScore', 'ef', 'probes']);
+    const { indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0, ef, probes } = params;
+
     const client = await this.pool.connect();
     try {
       const vectorStr = `[${queryVector.join(',')}]`;
@@ -78,13 +102,13 @@ export class PgVector extends MastraVector {
       // Set HNSW search parameter if applicable
       if (indexInfo.type === 'hnsw') {
         // Calculate ef and clamp between 1 and 1000
-        const calculatedEf = options?.ef ?? Math.max(topK, (indexInfo?.config?.m ?? 16) * topK);
+        const calculatedEf = ef ?? Math.max(topK, (indexInfo?.config?.m ?? 16) * topK);
         const searchEf = Math.min(1000, 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}`);
+      if (indexInfo.type === 'ivfflat' && probes) {
+        await client.query(`SET LOCAL ivfflat.probes = ${probes}`);
       }
 
       const query = `
@@ -115,12 +139,11 @@ export class PgVector extends MastraVector {
     }
   }
 
-  async upsert(
-    indexName: string,
-    vectors: number[][],
-    metadata?: Record<string, any>[],
-    ids?: string[],
-  ): Promise<string[]> {
+  async upsert(...args: ParamsToArgs<UpsertVectorParams>): Promise<string[]> {
+    const params = this.normalizeArgs<UpsertVectorParams>('upsert', args);
+
+    const { indexName, vectors, metadata, ids } = params;
+
     // Start a transaction
     const client = await this.pool.connect();
     try {
@@ -151,13 +174,11 @@ export class PgVector extends MastraVector {
     }
   }
 
-  async createIndex(
-    indexName: string,
-    dimension: number,
-    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
-    indexConfig: IndexConfig = {},
-    defineIndex: boolean = true,
-  ): Promise<void> {
+  async createIndex(...args: ParamsToArgs<PgCreateIndexParams>): Promise<void> {
+    const params = this.normalizeArgs<PgCreateIndexParams>('createIndex', args, ['indexConfig', 'buildIndex']);
+
+    const { indexName, dimension, metric = 'cosine', indexConfig = {}, buildIndex = true } = params;
+
     const client = await this.pool.connect();
     try {
       // Validate inputs
@@ -192,8 +213,8 @@ export class PgVector extends MastraVector {
         );
       `);
 
-      if (defineIndex) {
-        await this.defineIndex(indexName, metric, indexConfig);
+      if (buildIndex) {
+        await this.buildIndex({ indexName, metric, indexConfig });
       }
     } catch (error: any) {
       console.error('Failed to create vector table:', error);
@@ -203,11 +224,22 @@ export class PgVector extends MastraVector {
     }
   }
 
+  /**
+   * @deprecated This function is deprecated. Use buildIndex instead
+   */
   async defineIndex(
     indexName: string,
     metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
     indexConfig: IndexConfig,
   ): Promise<void> {
+    return this.buildIndex({ indexName, metric, indexConfig });
+  }
+
+  async buildIndex(...args: ParamsToArgs<PgDefineIndexParams>): Promise<void> {
+    const params = this.normalizeArgs<PgDefineIndexParams>('buildIndex', args, ['metric', 'indexConfig']);
+
+    const { indexName, metric = 'cosine', indexConfig } = params;
+
     const client = await this.pool.connect();
     try {
       await client.query(`DROP INDEX IF EXISTS ${indexName}_vector_idx`);

package/src/vector/performance.helpers.ts

@@ -1,6 +1,6 @@
-import { type IndexConfig, type IndexType } from './types';
+import type { IndexConfig, IndexType } from './types';
 
-import { PgVector } from '.';
+import type { PgVector } from '.';
 
 export interface TestResult {
   distribution: string;
@@ -236,7 +236,7 @@ export interface TestConfig {
 
 export async function warmupQuery(vectorDB: PgVector, indexName: string, dimension: number, k: number) {
   const warmupVector = generateRandomVectors(1, dimension)[0] as number[];
-  await vectorDB.query(indexName, warmupVector, k);
+  await vectorDB.query({ indexName, queryVector: warmupVector, topK: k });
 }
 
 export async function measureLatency<T>(fn: () => Promise<T>): Promise<[number, T]> {

package/src/vector/sql-builder.ts

@@ -1,12 +1,12 @@
-import {
-  type BasicOperator,
-  type NumericOperator,
-  type ArrayOperator,
-  type ElementOperator,
-  type LogicalOperator,
-  type RegexOperator,
-  type Filter,
-} from '@mastra/core/filter';
+import type {
+  BasicOperator,
+  NumericOperator,
+  ArrayOperator,
+  ElementOperator,
+  LogicalOperator,
+  RegexOperator,
+  VectorFilter,
+} from '@mastra/core/vector/filter';
 
 export type OperatorType =
   | BasicOperator
@@ -180,7 +180,7 @@ export const handleKey = (key: string) => {
   return key.replace(/\./g, ',');
 };
 
-export function buildFilterQuery(filter: Filter, minScore: number): FilterResult {
+export function buildFilterQuery(filter: VectorFilter, minScore: number): FilterResult {
   const values = [minScore];
 
   function buildCondition(key: string, value: any, parentPath: string): string {
@@ -232,7 +232,11 @@ export function buildFilterQuery(filter: Filter, minScore: number): FilterResult
     return operatorResult.sql;
   }
 
-  function handleLogicalOperator(key: '$and' | '$or' | '$not' | '$nor', value: Filter[], parentPath: string): string {
+  function handleLogicalOperator(
+    key: '$and' | '$or' | '$not' | '$nor',
+    value: VectorFilter[],
+    parentPath: string,
+  ): string {
     if (key === '$not') {
       // For top-level $not
       const entries = Object.entries(value);
@@ -256,8 +260,8 @@ export function buildFilterQuery(filter: Filter, minScore: number): FilterResult
     }
 
     const joinOperator = key === '$or' || key === '$nor' ? 'OR' : 'AND';
-    const conditions = value.map((f: Filter) => {
-      const entries = Object.entries(f);
+    const conditions = value.map((f: VectorFilter) => {
+      const entries = Object.entries(f || {});
       if (entries.length === 0) return '';
 
       const [firstKey, firstValue] = entries[0] || [];

package/src/vector/vector.performance.test.ts

@@ -1,10 +1,9 @@
 import pg from 'pg';
 import { describe, it, beforeAll, afterAll, beforeEach, afterEach } from 'vitest';
 
+import type { TestConfig, TestResult } from './performance.helpers';
 import {
   baseTestConfigs,
-  TestConfig,
-  TestResult,
   calculateTimeout,
   generateRandomVectors,
   findNearestBruteForce,
@@ -18,8 +17,9 @@ import {
   generateSkewedVectors,
   getHNSWConfig,
   getIndexDescription,
+  warmupQuery,
 } from './performance.helpers';
-import { IndexConfig, IndexType } from './types';
+import type { IndexConfig, IndexType } from './types';
 
 import { PgVector } from '.';
 
@@ -93,8 +93,7 @@ async function smartWarmup(
   const cacheKey = `${dimension}-${k}-${indexType}`;
   if (!warmupCache.has(cacheKey)) {
     console.log(`Warming up ${indexType} index for ${dimension}d vectors, k=${k}`);
-    const warmupVector = generateRandomVectors(1, dimension)[0] as number[];
-    await vectorDB.query(testIndexName, warmupVector, k);
+    await warmupQuery(vectorDB, testIndexName, dimension, k);
     warmupCache.set(cacheKey, true);
   }
 }
@@ -163,13 +162,13 @@ describe('PostgreSQL Index Performance', () => {
               // Create index and insert vectors
               const lists = getListCount(indexConfig, testConfig.size);
 
-              await vectorDB.createIndex(
-                testIndexName,
-                testConfig.dimension,
-                'cosine',
+              await vectorDB.createIndex({
+                indexName: testIndexName,
+                dimension: testConfig.dimension,
+                metric: 'cosine',
                 indexConfig,
-                indexType === 'ivfflat',
-              );
+                buildIndex: indexType === 'ivfflat',
+              });
 
               console.log(
                 `Batched bulk upserting ${testVectors.length} ${distType} vectors into index ${testIndexName}`,
@@ -178,7 +177,7 @@ describe('PostgreSQL Index Performance', () => {
               await batchedBulkUpsert(vectorDB, testIndexName, testVectors, batchSizes);
               if (indexType === 'hnsw' || rebuild) {
                 console.log('rebuilding index');
-                await vectorDB.defineIndex(testIndexName, 'cosine', indexConfig);
+                await vectorDB.buildIndex({ indexName: testIndexName, metric: 'cosine', indexConfig });
                 console.log('index rebuilt');
               }
               await smartWarmup(vectorDB, testIndexName, indexType, testConfig.dimension, testConfig.k);
@@ -194,15 +193,12 @@ describe('PostgreSQL Index Performance', () => {
                   const expectedNeighbors = findNearestBruteForce(queryVector, testVectors, testConfig.k);
 
                   const [latency, actualResults] = await measureLatency(async () =>
-                    vectorDB.query(
-                      testIndexName,
+                    vectorDB.query({
+                      indexName: testIndexName,
                       queryVector,
-                      testConfig.k,
-                      undefined,
-                      false,
-                      0,
-                      { ef }, // For HNSW
-                    ),
+                      topK: testConfig.k,
+                      ef, // For HNSW
+                    }),
                   );
 
                   const actualNeighbors = actualResults.map(r => r.metadata?.index);