@mastra/pg 0.1.0-alpha.10

package/src/vector/index.ts

@@ -0,0 +1,282 @@
+import type { Filter } from '@mastra/core/filter';
+import { type IndexStats, type QueryResult, MastraVector } from '@mastra/core/vector';
+import pg from 'pg';
+
+import { PGFilterTranslator } from './filter';
+import { buildFilterQuery } from './sql-builder';
+
+export class PgVector extends MastraVector {
+  private pool: pg.Pool;
+
+  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 query(
+    indexName: string,
+    queryVector: number[],
+    topK: number = 10,
+    filter?: Filter,
+    includeVector: boolean = false,
+    minScore: number = 0, // Optional minimum score threshold
+  ): 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);
+
+      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',
+  ): 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
+        );
+      `);
+
+      // Create the index
+      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);
+      `);
+    } catch (error: any) {
+      console.error('Failed to create vector table:', error);
+      throw error;
+    } 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<IndexStats> {
+    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 metricQuery = `
+            SELECT
+                am.amname as index_method,
+                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, metricResult] = await Promise.all([
+        client.query(dimensionQuery, [indexName]),
+        client.query(countQuery),
+        client.query(metricQuery),
+      ]);
+
+      // Convert pg_vector index method to our metric type
+      let metric: 'cosine' | 'euclidean' | 'dotproduct' = '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';
+        }
+      }
+
+      return {
+        dimension: dimResult.rows[0].dimension,
+        count: parseInt(countResult.rows[0].count),
+        metric,
+      };
+    } 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/sql-builder.ts

@@ -0,0 +1,285 @@
+import {
+  BasicOperator,
+  NumericOperator,
+  ArrayOperator,
+  ElementOperator,
+  LogicalOperator,
+  RegexOperator,
+  Filter,
+} from '@mastra/core/filter';
+
+export type OperatorType =
+  | BasicOperator
+  | NumericOperator
+  | ArrayOperator
+  | ElementOperator
+  | LogicalOperator
+  | '$contains'
+  | Exclude<RegexOperator, '$options'>;
+
+type FilterOperator = {
+  sql: string;
+  needsValue: boolean;
+  transformValue?: (value: any) => any;
+};
+
+type OperatorFn = (key: string, paramIndex: number, value?: any) => FilterOperator;
+
+// Helper functions to create operators
+const createBasicOperator = (symbol: string) => {
+  return (key: string, paramIndex: number) => ({
+    sql: `CASE 
+      WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${handleKey(key)}}' IS ${symbol === '=' ? '' : 'NOT'} NULL
+      ELSE metadata#>>'{${handleKey(key)}}' ${symbol} $${paramIndex}::text
+    END`,
+    needsValue: true,
+  });
+};
+
+const createNumericOperator = (symbol: string) => {
+  return (key: string, paramIndex: number) => ({
+    sql: `(metadata#>>'{${handleKey(key)}}')::numeric ${symbol} $${paramIndex}`,
+    needsValue: true,
+  });
+};
+
+function buildElemMatchConditions(value: any, paramIndex: number): { sql: string; values: any[] } {
+  if (typeof value !== 'object' || Array.isArray(value)) {
+    throw new Error('$elemMatch requires an object with conditions');
+  }
+
+  const conditions: string[] = [];
+  const values: any[] = [];
+
+  Object.entries(value).forEach(([field, val]) => {
+    const nextParamIndex = paramIndex + values.length;
+
+    let paramOperator;
+    let paramKey;
+    let paramValue;
+
+    if (field.startsWith('$')) {
+      paramOperator = field;
+      paramKey = '';
+      paramValue = val;
+    } else if (typeof val === 'object' && !Array.isArray(val)) {
+      const [op, opValue] = Object.entries(val || {})[0] || [];
+      paramOperator = op;
+      paramKey = field;
+      paramValue = opValue;
+    } else {
+      paramOperator = '$eq';
+      paramKey = field;
+      paramValue = val;
+    }
+
+    const operatorFn = FILTER_OPERATORS[paramOperator as keyof typeof FILTER_OPERATORS];
+    if (!operatorFn) {
+      throw new Error(`Invalid operator: ${paramOperator}`);
+    }
+    const result = operatorFn(paramKey, nextParamIndex, paramValue);
+
+    const sql = result.sql.replaceAll('metadata#>>', 'elem#>>');
+    conditions.push(sql);
+    if (result.needsValue) {
+      values.push(paramValue);
+    }
+  });
+
+  return {
+    sql: conditions.join(' AND '),
+    values,
+  };
+}
+
+// Define all filter operators
+export const FILTER_OPERATORS: Record<string, OperatorFn> = {
+  $eq: createBasicOperator('='),
+  $ne: createBasicOperator('!='),
+  $gt: createNumericOperator('>'),
+  $gte: createNumericOperator('>='),
+  $lt: createNumericOperator('<'),
+  $lte: createNumericOperator('<='),
+
+  // Array Operators
+  $in: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' = ANY($${paramIndex}::text[])`,
+    needsValue: true,
+  }),
+  $nin: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' != ALL($${paramIndex}::text[])`,
+    needsValue: true,
+  }),
+  $all: (key, paramIndex) => ({
+    sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false 
+          ELSE (metadata#>'{${handleKey(key)}}')::jsonb ?& $${paramIndex}::text[] END`,
+    needsValue: true,
+  }),
+  $elemMatch: (key: string, paramIndex: number, value: any): FilterOperator => {
+    const { sql, values } = buildElemMatchConditions(value, paramIndex);
+    return {
+      sql: `(
+        CASE
+          WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
+            EXISTS (
+              SELECT 1 
+              FROM jsonb_array_elements(metadata->'${handleKey(key)}') as elem
+              WHERE ${sql}
+            )
+          ELSE FALSE
+        END
+      )`,
+      needsValue: true,
+      transformValue: () => values,
+    };
+  },
+  // Element Operators
+  $exists: key => ({
+    sql: `metadata ? '${key}'`,
+    needsValue: false,
+  }),
+
+  // Logical Operators
+  $and: key => ({ sql: `(${key})`, needsValue: false }),
+  $or: key => ({ sql: `(${key})`, needsValue: false }),
+  $not: key => ({ sql: `NOT (${key})`, needsValue: false }),
+  $nor: key => ({ sql: `NOT (${key})`, needsValue: false }),
+
+  // Regex Operators
+  $regex: (key, paramIndex) => ({
+    sql: `metadata#>>'{${handleKey(key)}}' ~ $${paramIndex}`,
+    needsValue: true,
+  }),
+
+  $contains: (key, paramIndex) => ({
+    sql: `metadata @> $${paramIndex}::jsonb`,
+    needsValue: true,
+    transformValue: value => {
+      const parts = key.split('.');
+      return JSON.stringify(parts.reduceRight((value, key) => ({ [key]: value }), value));
+    },
+  }),
+  $size: (key: string, paramIndex: number) => ({
+    sql: `(
+      CASE
+        WHEN jsonb_typeof(metadata#>'{${handleKey(key)}}') = 'array' THEN 
+          jsonb_array_length(metadata#>'{${handleKey(key)}}') = $${paramIndex}
+        ELSE FALSE
+      END
+    )`,
+    needsValue: true,
+  }),
+};
+
+export interface FilterResult {
+  sql: string;
+  values: any[];
+}
+
+export const handleKey = (key: string) => {
+  return key.replace(/\./g, ',');
+};
+
+export function buildFilterQuery(filter: Filter, minScore: number): FilterResult {
+  const values = [minScore];
+
+  function buildCondition(key: string, value: any, parentPath: string): string {
+    // Handle logical operators ($and/$or)
+    if (['$and', '$or', '$not', '$nor'].includes(key)) {
+      return handleLogicalOperator(key as '$and' | '$or' | '$not' | '$nor', value, parentPath);
+    }
+
+    // If condition is not a FilterCondition object, assume it's an equality check
+    if (!value || typeof value !== 'object') {
+      values.push(value);
+      return `metadata#>>'{${handleKey(key)}}' = $${values.length}`;
+    }
+
+    // Handle operator conditions
+    const [[operator, operatorValue] = []] = Object.entries(value);
+
+    // Special handling for nested $not
+    if (operator === '$not') {
+      const entries = Object.entries(operatorValue as Record<string, unknown>);
+      const conditions = entries
+        .map(([nestedOp, nestedValue]) => {
+          if (!FILTER_OPERATORS[nestedOp as keyof typeof FILTER_OPERATORS]) {
+            throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
+          }
+          const operatorFn = FILTER_OPERATORS[nestedOp]!;
+          const operatorResult = operatorFn(key, values.length + 1);
+          if (operatorResult.needsValue) {
+            values.push(nestedValue as number);
+          }
+          return operatorResult.sql;
+        })
+        .join(' AND ');
+
+      return `NOT (${conditions})`;
+    }
+    const operatorFn = FILTER_OPERATORS[operator as string]!;
+    const operatorResult = operatorFn(key, values.length + 1, operatorValue);
+    if (operatorResult.needsValue) {
+      const transformedValue = operatorResult.transformValue
+        ? operatorResult.transformValue(operatorValue)
+        : operatorValue;
+      if (Array.isArray(transformedValue) && operator === '$elemMatch') {
+        values.push(...transformedValue);
+      } else {
+        values.push(transformedValue);
+      }
+    }
+    return operatorResult.sql;
+  }
+
+  function handleLogicalOperator(key: '$and' | '$or' | '$not' | '$nor', value: Filter[], parentPath: string): string {
+    if (key === '$not') {
+      // For top-level $not
+      const entries = Object.entries(value);
+      const conditions = entries
+        .map(([fieldKey, fieldValue]) => buildCondition(fieldKey, fieldValue, key))
+        .join(' AND ');
+      return `NOT (${conditions})`;
+    }
+
+    // Handle empty conditions
+    if (!value || value.length === 0) {
+      switch (key) {
+        case '$and':
+        case '$nor':
+          return 'true'; // Empty $and/$nor match everything
+        case '$or':
+          return 'false'; // Empty $or matches nothing
+        default:
+          return 'true';
+      }
+    }
+
+    const joinOperator = key === '$or' || key === '$nor' ? 'OR' : 'AND';
+    const conditions = value.map((f: Filter) => {
+      const entries = Object.entries(f);
+      if (entries.length === 0) return '';
+
+      const [firstKey, firstValue] = entries[0] || [];
+      if (['$and', '$or', '$not', '$nor'].includes(firstKey as string)) {
+        return buildCondition(firstKey as string, firstValue, parentPath);
+      }
+      return entries.map(([k, v]) => buildCondition(k, v, parentPath)).join(` ${joinOperator} `);
+    });
+
+    const joined = conditions.join(` ${joinOperator} `);
+    const operatorFn = FILTER_OPERATORS[key]!;
+    return operatorFn(joined, 0, value).sql;
+  }
+
+  if (!filter) {
+    return { sql: '', values };
+  }
+
+  const conditions = Object.entries(filter)
+    .map(([key, value]) => buildCondition(key, value, ''))
+    .filter(Boolean)
+    .join(' AND ');
+
+  return { sql: conditions ? `WHERE ${conditions}` : '', values };
+}

package/tsconfig.json

@@ -0,0 +1,15 @@
+{
+  "extends": "../../tsconfig.node.json",
+  "compilerOptions": {
+    "moduleResolution": "bundler",
+    "outDir": "./dist",
+    "rootDir": "./src",
+    "module": "ES2022",
+    "target": "ES2020",
+    "declaration": true,
+    "declarationMap": true,
+    "noEmit": false
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "**/*.test.ts"]
+}

package/vitest.config.ts

@@ -0,0 +1,11 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['src/**/*.test.ts'],
+    coverage: {
+      reporter: ['text', 'json', 'html'],
+    },
+  },
+});