@mastra/pg 0.0.0-commonjs-20250227130920

package/src/vector/index.ts

@@ -0,0 +1,391 @@
+import type { Filter } from '@mastra/core/filter';
+import { MastraVector } from '@mastra/core/vector';
+import type { IndexStats, QueryResult } from '@mastra/core/vector';
+import pg from 'pg';
+
+import { PGFilterTranslator } from './filter';
+import { buildFilterQuery } from './sql-builder';
+import type { IndexConfig, IndexType } from './types';
+
+export interface PGIndexStats extends IndexStats {
+  type: IndexType;
+  config: {
+    m?: number;
+    efConstruction?: number;
+    lists?: number;
+    probes?: number;
+  };
+}
+
+export class PgVector extends MastraVector {
+  private pool: pg.Pool;
+  private indexCache: Map<string, PGIndexStats> = new Map();
+
+  constructor(connectionString: string) {
+    super();
+
+    const basePool = new pg.Pool({
+      connectionString,
+      max: 20, // Maximum number of clients in the pool
+      idleTimeoutMillis: 30000, // Close idle connections after 30 seconds
+      connectionTimeoutMillis: 2000, // Fail fast if can't connect
+    });
+
+    const telemetry = this.__getTelemetry();
+
+    this.pool =
+      telemetry?.traceClass(basePool, {
+        spanNamePrefix: 'pg-vector',
+        attributes: {
+          'vector.type': 'postgres',
+        },
+      }) ?? basePool;
+  }
+
+  transformFilter(filter?: Filter) {
+    const pgFilter = new PGFilterTranslator();
+    const translatedFilter = pgFilter.translate(filter ?? {});
+    return translatedFilter;
+  }
+
+  async getIndexInfo(indexName: string): Promise<PGIndexStats> {
+    if (!this.indexCache.has(indexName)) {
+      this.indexCache.set(indexName, await this.describeIndex(indexName));
+    }
+    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[]> {
+    const client = await this.pool.connect();
+    try {
+      const vectorStr = `[${queryVector.join(',')}]`;
+      const translatedFilter = this.transformFilter(filter);
+      const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
+
+      // Get index type and configuration
+      const indexInfo = await this.getIndexInfo(indexName);
+
+      // 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 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}`);
+      }
+
+      const query = `
+        WITH vector_scores AS (
+          SELECT
+            vector_id as id,
+            1 - (embedding <=> '${vectorStr}'::vector) as score,
+            metadata
+            ${includeVector ? ', embedding' : ''}
+          FROM ${indexName}
+          ${filterQuery}
+        )
+        SELECT *
+        FROM vector_scores
+        WHERE score > $1
+        ORDER BY score DESC
+        LIMIT ${topK}`;
+      const result = await client.query(query, filterValues);
+
+      return result.rows.map(({ id, score, metadata, embedding }) => ({
+        id,
+        score,
+        metadata,
+        ...(includeVector && embedding && { vector: JSON.parse(embedding) }),
+      }));
+    } finally {
+      client.release();
+    }
+  }
+
+  async upsert(
+    indexName: string,
+    vectors: number[][],
+    metadata?: Record<string, any>[],
+    ids?: string[],
+  ): Promise<string[]> {
+    // Start a transaction
+    const client = await this.pool.connect();
+    try {
+      await client.query('BEGIN');
+      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
+        `;
+
+        await client.query(query, [vectorIds[i], `[${vectors[i]?.join(',')}]`, JSON.stringify(metadata?.[i] || {})]);
+      }
+
+      await client.query('COMMIT');
+      return vectorIds;
+    } catch (error) {
+      await client.query('ROLLBACK');
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  async createIndex(
+    indexName: string,
+    dimension: number,
+    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
+    indexConfig: IndexConfig = {},
+    defineIndex: boolean = true,
+  ): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      // Validate inputs
+      if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
+        throw new Error('Invalid index name format');
+      }
+      if (!Number.isInteger(dimension) || dimension <= 0) {
+        throw new Error('Dimension must be a positive integer');
+      }
+
+      // First check if vector extension is available
+      const extensionCheck = await client.query(`
+        SELECT EXISTS (
+          SELECT 1 FROM pg_available_extensions WHERE name = 'vector'
+        );
+      `);
+
+      if (!extensionCheck.rows[0].exists) {
+        throw new Error('PostgreSQL vector extension is not available. Please install it first.');
+      }
+
+      // Try to create extension
+      await client.query('CREATE EXTENSION IF NOT EXISTS vector');
+
+      // Create the table with explicit schema
+      await client.query(`
+        CREATE TABLE IF NOT EXISTS ${indexName} (
+          id SERIAL PRIMARY KEY,
+          vector_id TEXT UNIQUE NOT NULL,
+          embedding vector(${dimension}),
+          metadata JSONB DEFAULT '{}'::jsonb
+        );
+      `);
+
+      if (defineIndex) {
+        await this.defineIndex(indexName, metric, indexConfig);
+      }
+    } catch (error: any) {
+      console.error('Failed to create vector table:', error);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * @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(
+    indexName: string,
+    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
+    indexConfig: IndexConfig,
+  ): Promise<void> {
+    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: string;
+      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: number;
+        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(4000, 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(): Promise<string[]> {
+    const client = await this.pool.connect();
+    try {
+      // Then let's see which ones have vector columns
+      const vectorTablesQuery = `
+            SELECT DISTINCT table_name
+            FROM information_schema.columns
+            WHERE table_schema = 'public'
+            AND udt_name = 'vector';
+        `;
+      const vectorTables = await client.query(vectorTablesQuery);
+      return vectorTables.rows.map(row => row.table_name);
+    } finally {
+      client.release();
+    }
+  }
+
+  async describeIndex(indexName: string): Promise<PGIndexStats> {
+    const client = await this.pool.connect();
+    try {
+      // Get vector dimension
+      const dimensionQuery = `
+                SELECT atttypmod as dimension
+                FROM pg_attribute
+                WHERE attrelid = $1::regclass
+                AND attname = 'embedding';
+            `;
+
+      // Get row count
+      const countQuery = `
+                SELECT COUNT(*) as count
+                FROM ${indexName};
+            `;
+
+      // Get index metric type
+      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
+            JOIN pg_am am ON c.relam = am.oid
+            JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
+            WHERE c.relname = '${indexName}_vector_idx';
+            `;
+
+      const [dimResult, countResult, indexResult] = await Promise.all([
+        client.query(dimensionQuery, [indexName]),
+        client.query(countQuery),
+        client.query(indexQuery),
+      ]);
+
+      const { index_method, index_def, operator_class } = indexResult.rows[0] || {
+        index_method: 'flat',
+        index_def: '',
+        operator_class: 'cosine',
+      };
+
+      // Convert pg_vector index method to our metric type
+      const metric = operator_class.includes('l2')
+        ? 'euclidean'
+        : operator_class.includes('ip')
+          ? 'dotproduct'
+          : 'cosine';
+
+      // Parse index configuration
+      const config: { m?: number; efConstruction?: number; lists?: number } = {};
+
+      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,
+        type: index_method as 'flat' | 'hnsw' | 'ivfflat',
+        config,
+      };
+    } catch (e: any) {
+      await client.query('ROLLBACK');
+      throw new Error(`Failed to describe vector table: ${e.message}`);
+    } finally {
+      client.release();
+    }
+  }
+
+  async deleteIndex(indexName: string): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      // Drop the table
+      await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
+    } catch (error: any) {
+      await client.query('ROLLBACK');
+      throw new Error(`Failed to delete vector table: ${error.message}`);
+    } finally {
+      client.release();
+    }
+  }
+
+  async truncateIndex(indexName: string) {
+    const client = await this.pool.connect();
+    try {
+      await client.query(`TRUNCATE ${indexName}`);
+    } catch (e: any) {
+      await client.query('ROLLBACK');
+      throw new Error(`Failed to truncate vector table: ${e.message}`);
+    } finally {
+      client.release();
+    }
+  }
+
+  async disconnect() {
+    await this.pool.end();
+  }
+}

package/src/vector/performance.helpers.ts

@@ -0,0 +1,286 @@
+import type { IndexConfig, IndexType } from './types';
+
+import type { PgVector } from '.';
+
+export 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 const generateRandomVectors = (count: number, dim: number) => {
+  return Array.from({ length: count }, () => {
+    return Array.from({ length: dim }, () => Math.random() * 2 - 1);
+  });
+};
+
+export const generateClusteredVectors = (count: number, dim: number, numClusters: number = 10) => {
+  // Generate cluster centers
+  const centers = Array.from({ length: numClusters }, () => Array.from({ length: dim }, () => Math.random() * 2 - 1));
+
+  // Generate vectors around centers with varying spread
+  return Array.from({ length: count }, () => {
+    // Pick a random cluster, with some clusters being more popular
+    const centerIdx = Math.floor(Math.pow(Math.random(), 2) * numClusters);
+    const center = centers[centerIdx] as number[];
+
+    // Add noise, with some vectors being further from centers
+    const spread = Math.random() < 0.8 ? 0.1 : 0.5; // 80% close, 20% far
+    return center.map(c => c + (Math.random() * spread - spread / 2));
+  });
+};
+
+// Or even more extreme:
+export const generateSkewedVectors = (count: number, dim: number) => {
+  // Create dense clusters with sparse regions
+  const vectors: number[][] = [];
+
+  const denseCount = Math.floor(count * 0.6);
+  const sparseCount = count - denseCount;
+
+  // Dense cluster (60% of vectors)
+  const denseCenter = Array.from({ length: dim }, () => Math.random() * 0.2);
+  for (let i = 0; i < denseCount; i++) {
+    vectors.push(denseCenter.map(c => c + (Math.random() * 0.1 - 0.05)));
+  }
+
+  // Scattered vectors (40%)
+  for (let i = 0; i < sparseCount; i++) {
+    vectors.push(Array.from({ length: dim }, () => Math.random() * 2 - 1));
+  }
+
+  return vectors.sort(() => Math.random() - 0.5); // Shuffle
+};
+
+export const findNearestBruteForce = (query: number[], vectors: number[][], k: number) => {
+  const similarities = vectors.map((vector, idx) => {
+    const similarity = cosineSimilarity(query, vector);
+    return { idx, dist: similarity };
+  });
+
+  const sorted = similarities.sort((a, b) => b.dist - a.dist);
+  return sorted.slice(0, k).map(x => x.idx);
+};
+
+export const calculateRecall = (actual: number[], expected: number[], k: number): number => {
+  let score = 0;
+  for (let i = 0; i < k; i++) {
+    if (actual[i] === expected[i]) {
+      score += 1;
+    } else if (expected.includes(actual[i] ?? 0)) {
+      score += 0.5;
+    }
+  }
+  return score / k;
+};
+
+export function cosineSimilarity(a: number[], b: number[]): number {
+  const dotProduct = a.reduce((sum, val, i) => sum + (val ?? 0) * (b[i] ?? 0), 0);
+  const normA = Math.sqrt(a.reduce((sum, val) => sum + val * val, 0));
+  const normB = Math.sqrt(b.reduce((sum, val) => sum + val * val, 0));
+  return dotProduct / (normA * normB);
+}
+
+export const formatTable = (data: any[], columns: string[]) => {
+  const colWidths = columns.map(col =>
+    Math.max(
+      col.length,
+      ...data.map(row => {
+        const value = row[col];
+        return value === undefined || value === null ? '-'.length : value.toString().length;
+      }),
+    ),
+  );
+
+  const topBorder = '┌' + colWidths.map(w => '─'.repeat(w)).join('┬') + '┐';
+  const headerSeparator = '├' + colWidths.map(w => '─'.repeat(w)).join('┼') + '┤';
+  const bottomBorder = '└' + colWidths.map(w => '─'.repeat(w)).join('┴') + '┘';
+
+  const header = '│' + columns.map((col, i) => col.padEnd(colWidths[i] ?? 0)).join('│') + '│';
+  const rows = data.map(
+    row =>
+      '│' +
+      columns
+        .map((col, i) => {
+          const value = row[col];
+          const displayValue = value === undefined || value === null ? '-' : value.toString();
+          return displayValue.padEnd(colWidths[i]);
+        })
+        .join('│') +
+      '│',
+  );
+
+  return [topBorder, header, headerSeparator, ...rows, bottomBorder].join('\n');
+};
+
+export const groupBy = <T, K extends keyof T>(
+  array: T[],
+  key: K | ((item: T) => string),
+  reducer?: (group: T[]) => any,
+): Record<string, any> => {
+  const grouped = array.reduce(
+    (acc, item) => {
+      const value = typeof key === 'function' ? key(item) : item[key];
+      if (!acc[value as any]) acc[value as any] = [];
+      acc[value as any]?.push(item);
+      return acc;
+    },
+    {} as Record<string, T[]>,
+  );
+
+  if (reducer) {
+    return Object.fromEntries(Object.entries(grouped).map(([key, group]) => [key, reducer(group)]));
+  }
+
+  return grouped;
+};
+
+export const calculateTimeout = (dimension: number, size: number, k: number) => {
+  let timeout = 600000;
+  if (dimension >= 1024) timeout *= 3;
+  else if (dimension >= 384) timeout *= 1.5;
+  if (size >= 10000) timeout *= 2;
+  if (k >= 75) timeout *= 1.5;
+  return timeout * 5;
+};
+
+export const baseTestConfigs = {
+  smokeTests: [{ dimension: 384, size: 1_000, k: 10, queryCount: 10 }],
+  '64': [
+    { dimension: 64, size: 100, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100, k: 25, queryCount: 30 },
+    { dimension: 64, size: 100, k: 50, queryCount: 30 },
+    { dimension: 64, size: 100, k: 100, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 64, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 25, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 50, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 100, queryCount: 30 },
+    { dimension: 64, size: 500_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 1_000_000, k: 10, queryCount: 30 },
+  ],
+  '384': [
+    { dimension: 384, size: 100, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100, k: 25, queryCount: 30 },
+    { dimension: 384, size: 100, k: 50, queryCount: 30 },
+    { dimension: 384, size: 100, k: 100, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 384, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 25, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 50, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 100, queryCount: 30 },
+    { dimension: 384, size: 500_000, k: 10, queryCount: 30 },
+  ],
+  '1024': [
+    { dimension: 1024, size: 100, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 50_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 50_000, k: 25, queryCount: 30 },
+  ],
+  stressTests: [
+    // Maximum load
+    { dimension: 512, size: 1_000_000, k: 50, queryCount: 5 },
+
+    // Dense search
+    { dimension: 256, size: 1_000_000, k: 100, queryCount: 5 },
+
+    { dimension: 1024, size: 500_000, k: 50, queryCount: 5 },
+  ],
+};
+
+export interface TestConfig {
+  dimension: number;
+  size: number;
+  k: number;
+  queryCount: number;
+}
+
+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);
+}
+
+export async function measureLatency<T>(fn: () => Promise<T>): Promise<[number, T]> {
+  const start = performance.now();
+  const result = await fn();
+  const end = performance.now();
+  return [end - start, result];
+}
+
+export const getListCount = (indexConfig: IndexConfig, size: number): number | undefined => {
+  if (indexConfig.type !== 'ivfflat') return undefined;
+  if (indexConfig.ivf?.lists) return indexConfig.ivf.lists;
+  return Math.max(100, Math.min(4000, Math.floor(Math.sqrt(size) * 2)));
+};
+
+export const getHNSWConfig = (indexConfig: IndexConfig): { m: number; efConstruction: number } => {
+  return {
+    m: indexConfig.hnsw?.m ?? 8,
+    efConstruction: indexConfig.hnsw?.efConstruction ?? 32,
+  };
+};
+
+export function getSearchEf(k: number, m: number) {
+  return {
+    default: Math.max(k, m * k), // Default calculation
+    lower: Math.max(k, (m * k) / 2), // Lower quality, faster
+    higher: Math.max(k, m * k * 2), // Higher quality, slower
+  };
+}
+
+export function getIndexDescription({
+  type,
+  hnsw,
+}: {
+  type: IndexType;
+  hnsw: { m: number; efConstruction: number };
+}): string {
+  if (type === 'hnsw') {
+    return `HNSW(m=${hnsw.m},ef=${hnsw.efConstruction})`;
+  }
+
+  if (type === 'ivfflat') {
+    return `IVF`;
+  }
+
+  return 'Flat';
+}