@mastra/rag 0.0.2-alpha.0

package/src/pg/index.ts

@@ -0,0 +1,255 @@
+import { IndexStats, QueryResult, MastraVector } from '@mastra/core';
+import { Pool } from 'pg';
+
+export class PgVector extends MastraVector {
+  private pool: Pool;
+
+  constructor(connectionString: string) {
+    super();
+
+    const basePool = new 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;
+  }
+
+  async query(
+    indexName: string,
+    queryVector: number[],
+    topK: number = 10,
+    filter?: Record<string, any>,
+    minScore: number = 0, // Optional minimum score threshold
+  ): Promise<QueryResult[]> {
+    const client = await this.pool.connect();
+    try {
+      let filterQuery = '';
+      let filterValues: any[] = [minScore];
+      const vectorStr = `[${queryVector.join(',')}]`;
+
+      if (filter) {
+        const conditions = Object.entries(filter).map(([key, value], index) => {
+          filterValues.push(value);
+          return `metadata->>'${key}' = $${index + 2}`; // +2 because $1 is minScore
+        });
+        if (conditions.length > 0) {
+          filterQuery = 'AND ' + conditions.join(' AND ');
+        }
+      }
+
+      const query = `
+            WITH vector_scores AS (
+                SELECT
+                    vector_id as id,
+                    1 - (embedding <=> '${vectorStr}'::vector) as score,
+                    metadata
+                FROM ${indexName}
+                WHERE true ${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(row => ({
+        id: row.id,
+        score: row.score,
+        metadata: row.metadata,
+      }));
+    } 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] || {})]);
+      }
+      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 {
+      // Create the extension if it doesn't exist
+      await client.query('CREATE EXTENSION IF NOT EXISTS vector');
+
+      // Create the table
+      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 an index for vector similarity search based on metric
+      const indexMethod =
+        metric === 'cosine' ? 'vector_cosine_ops' : metric === 'euclidean' ? 'vector_l2_ops' : 'vector_ip_ops'; // for dotproduct
+
+      await client.query(`
+                CREATE INDEX IF NOT EXISTS ${indexName}_vector_idx
+                ON ${indexName}
+                USING ivfflat (embedding ${indexMethod})
+                WITH (lists = 100);
+            `);
+    } catch (error: any) {
+      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/pg/index_test.ts

@@ -0,0 +1,212 @@
+import { describe, it, expect, beforeAll, afterAll, beforeEach, afterEach } from '@jest/globals';
+
+import { PgVector } from '.';
+
+describe('PgVector', () => {
+  let pgVector: PgVector;
+  const testIndexName = 'test_vectors';
+  const testIndexName2 = 'test_vectors1';
+  const connectionString = process.env.DB_URL || 'postgresql://postgres:postgres@localhost:5433/mastra';
+
+  beforeAll(async () => {
+    // Initialize PgVector
+    pgVector = new PgVector(connectionString);
+  });
+
+  afterAll(async () => {
+    // Clean up test tables
+
+    await pgVector.deleteIndex(testIndexName);
+
+    await pgVector.disconnect();
+  });
+
+  describe('createIndex', () => {
+    afterAll(async () => {
+      await pgVector.deleteIndex(testIndexName2);
+    });
+
+    it('should create a new vector table with specified dimensions', async () => {
+      await pgVector.createIndex(testIndexName, 3);
+
+      const stats = await pgVector.describeIndex(testIndexName);
+      expect(stats?.dimension).toBe(3);
+      expect(stats?.count).toBe(0);
+    });
+
+    it('should create index with specified metric', async () => {
+      await pgVector.createIndex(testIndexName2, 3, 'euclidean');
+
+      const stats = await pgVector.describeIndex(testIndexName2);
+
+      expect(stats.metric).toBe('euclidean');
+    });
+
+    it('should throw error if dimension is invalid', async () => {
+      await expect(pgVector.createIndex(`testIndexNameFail`, 0)).rejects.toThrow();
+    });
+  });
+
+  describe('upsert', () => {
+    beforeEach(async () => {
+      await pgVector.createIndex(testIndexName, 3);
+    });
+
+    afterEach(async () => {
+      await pgVector.deleteIndex(testIndexName);
+    });
+
+    it('should insert new vectors', async () => {
+      const vectors = [
+        [1, 2, 3],
+        [4, 5, 6],
+      ];
+      const ids = await pgVector.upsert(testIndexName, vectors);
+
+      expect(ids).toHaveLength(2);
+      const stats = await pgVector.describeIndex(testIndexName);
+      expect(stats.count).toBe(2);
+    });
+
+    it('should update existing vectors', async () => {
+      const vectors = [[1, 2, 3]];
+      const metadata = [{ test: 'initial' }];
+      const [id] = await pgVector.upsert(testIndexName, vectors, metadata);
+
+      // Update the same vector
+      const updatedVectors = [[4, 5, 6]];
+      const updatedMetadata = [{ test: 'updated' }];
+      await pgVector.upsert(testIndexName, updatedVectors, updatedMetadata, [id!]);
+
+      const results = await pgVector.query(testIndexName, [4, 5, 6], 1);
+      expect(results[0]?.id).toBe(id);
+      expect(results[0]?.metadata).toEqual({ test: 'updated' });
+    });
+
+    it('should handle metadata correctly', async () => {
+      const vectors = [[1, 2, 3]];
+      const metadata = [{ test: 'value', num: 123 }];
+
+      await pgVector.upsert(testIndexName, vectors, metadata);
+      const results = await pgVector.query(testIndexName, [1, 2, 3], 1);
+
+      expect(results[0]?.metadata).toEqual(metadata[0]);
+    });
+
+    it('should throw error if vector dimensions dont match', async () => {
+      const vectors = [[1, 2, 3, 4]]; // 4D vector for 3D index
+      await expect(pgVector.upsert(testIndexName, vectors)).rejects.toThrow();
+    });
+  });
+
+  describe('query', () => {
+    const indexName = 'test_query_2';
+    beforeAll(async () => {
+      // Drop if exists first
+      try {
+        await pgVector.deleteIndex(indexName);
+      } catch (e) {
+        // Ignore if doesn't exist
+      }
+
+      // Create fresh
+      await pgVector.createIndex(indexName, 3);
+    });
+
+    beforeEach(async () => {
+      // Clear the table first
+      await pgVector.truncateIndex(indexName);
+
+      const vectors = [
+        [1, 0, 0],
+        [0.8, 0.2, 0],
+        [0, 1, 0],
+      ];
+      const metadata = [
+        { type: 'a', value: 1 },
+        { type: 'b', value: 2 },
+        { type: 'a', value: 3 },
+      ];
+      await pgVector.upsert(indexName, vectors, metadata);
+    });
+
+    afterAll(async () => {
+      console.log('deleting index');
+      console.log(await pgVector.listIndexes());
+      await pgVector.deleteIndex(indexName);
+      console.log(await pgVector.listIndexes());
+    });
+
+    it('should return closest vectors', async () => {
+      const results = await pgVector.query(indexName, [1, 0, 0], 1);
+      expect(results).toHaveLength(1);
+      expect(results[0]?.score).toBeCloseTo(1, 5);
+    });
+
+    it('should respect topK parameter', async () => {
+      const results = await pgVector.query(indexName, [1, 0, 0], 2);
+      expect(results).toHaveLength(2);
+    });
+
+    it('should handle filters correctly', async () => {
+      const results = await pgVector.query(indexName, [1, 0, 0], 10, { type: 'a' });
+
+      expect(results).toHaveLength(1);
+      results.forEach(result => {
+        expect(result?.metadata?.type).toBe('a');
+      });
+    });
+  });
+
+  describe('listIndexes', () => {
+    const indexName = 'test_query_3';
+    beforeAll(async () => {
+      await pgVector.createIndex(indexName, 3);
+    });
+
+    afterAll(async () => {
+      await pgVector.deleteIndex(indexName);
+    });
+
+    it('should list all vector tables', async () => {
+      const indexes = await pgVector.listIndexes();
+      expect(indexes).toContain(indexName);
+    });
+
+    it('should return empty array when no indexes exist', async () => {
+      await pgVector.deleteIndex(indexName);
+      const indexes = await pgVector.listIndexes();
+      expect(indexes).toEqual([]);
+    });
+  });
+
+  describe('describeIndex', () => {
+    const indexName = 'test_query_4';
+    beforeAll(async () => {
+      await pgVector.createIndex(indexName, 3);
+    });
+
+    afterAll(async () => {
+      await pgVector.deleteIndex(indexName);
+    });
+    it('should return correct index stats', async () => {
+      await pgVector.createIndex(indexName, 3, 'cosine');
+      const vectors = [
+        [1, 2, 3],
+        [4, 5, 6],
+      ];
+      await pgVector.upsert(indexName, vectors);
+
+      const stats = await pgVector.describeIndex(indexName);
+      expect(stats).toEqual({
+        dimension: 3,
+        count: 2,
+        metric: 'cosine',
+      });
+    });
+
+    it('should throw error for non-existent index', async () => {
+      await expect(pgVector.describeIndex('non_existent')).rejects.toThrow();
+    });
+  });
+});

package/src/pinecone/index.test.ts

@@ -0,0 +1,130 @@
+import { describe, it, expect } from '@jest/globals';
+import dotenv from 'dotenv';
+
+import { PineconeVector } from './';
+
+dotenv.config();
+
+const PINECONE_API_KEY = process.env.PINECONE_API_KEY!;
+
+// if (!PINECONE_API_KEY) {
+//   throw new Error('Please set PINECONE_API_KEY and PINECONE_ENVIRONMENT in .env file');
+// }
+// TODO: skip until we the secrets on Github
+describe.skip('PineconeVector Integration Tests', () => {
+  let pineconeVector: PineconeVector;
+  const testIndexName = 'test-index-' + Date.now(); // Unique index name for each test run
+  const dimension = 3;
+
+  beforeAll(async () => {
+    pineconeVector = new PineconeVector(PINECONE_API_KEY);
+    // Create test index
+    await pineconeVector.createIndex(testIndexName, dimension);
+    // Wait for index to be ready
+    await new Promise(resolve => setTimeout(resolve, 60000));
+  }, 500000);
+
+  afterAll(async () => {
+    // Cleanup: delete test index
+    await pineconeVector.deleteIndex(testIndexName);
+  }, 500000);
+
+  describe('Index Operations', () => {
+    it('should list indexes including our test index', async () => {
+      const indexes = await pineconeVector.listIndexes();
+      expect(indexes).toContain(testIndexName);
+    }, 500000);
+
+    it('should describe index with correct properties', async () => {
+      const stats = await pineconeVector.describeIndex(testIndexName);
+      expect(stats.dimension).toBe(dimension);
+      expect(stats.metric).toBe('cosine');
+      expect(typeof stats.count).toBe('number');
+    }, 500000);
+  });
+
+  describe('Vector Operations', () => {
+    const testVectors = [
+      [1.0, 0.0, 0.0],
+      [0.0, 1.0, 0.0],
+      [0.0, 0.0, 1.0],
+    ];
+    const testMetadata = [{ label: 'x-axis' }, { label: 'y-axis' }, { label: 'z-axis' }];
+    let vectorIds: string[];
+
+    it('should upsert vectors with metadata', async () => {
+      vectorIds = await pineconeVector.upsert(testIndexName, testVectors, testMetadata);
+      expect(vectorIds).toHaveLength(3);
+      // Wait for vectors to be indexed
+      await new Promise(resolve => setTimeout(resolve, 5000));
+    }, 500000);
+
+    it.skip('should query vectors and return nearest neighbors', async () => {
+      const queryVector = [1.0, 0.1, 0.1];
+      const results = await pineconeVector.query(testIndexName, queryVector, 3);
+
+      expect(results).toHaveLength(3);
+      expect(results[0]!.score).toBeGreaterThan(0);
+      expect(results[0]!.metadata).toBeDefined();
+    }, 500000);
+
+    it.skip('should query vectors with metadata filter', async () => {
+      const queryVector = [0.0, 1.0, 0.0];
+      const filter = { label: 'y-axis' };
+
+      const results = await pineconeVector.query(testIndexName, queryVector, 1, filter);
+
+      expect(results).toHaveLength(1);
+      expect(results?.[0]?.metadata?.label).toBe('y-axis');
+    }, 500000);
+  });
+
+  describe('Error Handling', () => {
+    it('should handle non-existent index query gracefully', async () => {
+      const nonExistentIndex = 'non-existent-index';
+      await expect(pineconeVector.query(nonExistentIndex, [1, 0, 0])).rejects.toThrow();
+    }, 500000);
+
+    it('should handle incorrect dimension vectors', async () => {
+      const wrongDimVector = [[1, 0]]; // 2D vector for 3D index
+      await expect(pineconeVector.upsert(testIndexName, wrongDimVector)).rejects.toThrow();
+    }, 500000);
+  });
+
+  describe('Performance Tests', () => {
+    it('should handle batch upsert of 1000 vectors', async () => {
+      const batchSize = 1000;
+      const vectors = Array(batchSize)
+        .fill(null)
+        .map(() =>
+          Array(dimension)
+            .fill(null)
+            .map(() => Math.random()),
+        );
+      const metadata = vectors.map((_, i) => ({ id: i }));
+
+      const start = Date.now();
+      const ids = await pineconeVector.upsert(testIndexName, vectors, metadata);
+      const duration = Date.now() - start;
+
+      expect(ids).toHaveLength(batchSize);
+      console.log(`Batch upsert of ${batchSize} vectors took ${duration}ms`);
+    }, 300000); // 5 minute timeout
+
+    it('should perform multiple concurrent queries', async () => {
+      const queryVector = [1, 0, 0];
+      const numQueries = 10;
+
+      const start = Date.now();
+      const promises = Array(numQueries)
+        .fill(null)
+        .map(() => pineconeVector.query(testIndexName, queryVector));
+
+      const results = await Promise.all(promises);
+      const duration = Date.now() - start;
+
+      expect(results).toHaveLength(numQueries);
+      console.log(`${numQueries} concurrent queries took ${duration}ms`);
+    }, 500000);
+  });
+});

package/src/pinecone/index.ts

@@ -0,0 +1,118 @@
+import { MastraVector, QueryResult, IndexStats } from '@mastra/core';
+import { Pinecone } from '@pinecone-database/pinecone';
+
+export class PineconeVector extends MastraVector {
+  private client: Pinecone;
+
+  constructor(apiKey: string, environment?: string) {
+    super();
+
+    const opts: { apiKey: string; controllerHostUrl?: string } = { apiKey };
+
+    if (environment) {
+      opts['controllerHostUrl'] = environment;
+    }
+
+    const baseClient = new Pinecone(opts);
+    const telemetry = this.__getTelemetry();
+    this.client =
+      telemetry?.traceClass(baseClient, {
+        spanNamePrefix: 'pinecone-vector',
+        attributes: {
+          'vector.type': 'pinecone',
+        },
+      }) ?? baseClient;
+  }
+
+  async createIndex(
+    indexName: string,
+    dimension: number,
+    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
+  ): Promise<void> {
+    await this.client.createIndex({
+      name: indexName,
+      dimension: dimension,
+      metric: metric,
+      spec: {
+        serverless: {
+          cloud: 'aws',
+          region: 'us-east-1',
+        },
+      },
+    });
+  }
+
+  async upsert(
+    indexName: string,
+    vectors: number[][],
+    metadata?: Record<string, any>[],
+    ids?: string[],
+  ): Promise<string[]> {
+    const index = this.client.Index(indexName);
+
+    // Generate IDs if not provided
+    const vectorIds = ids || vectors.map(() => crypto.randomUUID());
+
+    const records = vectors.map((vector, i) => ({
+      id: vectorIds[i]!,
+      values: vector,
+      metadata: metadata?.[i] || {},
+    }));
+
+    // Pinecone has a limit of 100 vectors per upsert request
+    const batchSize = 100;
+    for (let i = 0; i < records.length; i += batchSize) {
+      const batch = records.slice(i, i + batchSize);
+      await index.upsert(batch);
+    }
+
+    return vectorIds;
+  }
+
+  async query(
+    indexName: string,
+    queryVector: number[],
+    topK: number = 10,
+    filter?: Record<string, any>,
+  ): Promise<QueryResult[]> {
+    const index = this.client.Index(indexName);
+
+    const results = await index.query({
+      vector: queryVector,
+      topK,
+      filter,
+      includeMetadata: true,
+    });
+
+    return results.matches.map(match => ({
+      id: match.id,
+      score: match.score || 0,
+      metadata: match.metadata as Record<string, any>,
+    }));
+  }
+
+  async listIndexes(): Promise<string[]> {
+    const indexesResult = await this.client.listIndexes();
+    return indexesResult?.indexes?.map(index => index.name) || [];
+  }
+
+  async describeIndex(indexName: string): Promise<IndexStats> {
+    const index = this.client.Index(indexName);
+    const stats = await index.describeIndexStats();
+    const description = await this.client.describeIndex(indexName);
+
+    return {
+      dimension: description.dimension,
+      count: stats.totalRecordCount || 0,
+      metric: description.metric as 'cosine' | 'euclidean' | 'dotproduct',
+    };
+  }
+
+  async deleteIndex(indexName: string): Promise<void> {
+    try {
+      await this.client.deleteIndex(indexName);
+    } catch (error: any) {
+      throw new Error(`Failed to delete Pinecone index: ${error.message}`);
+    }
+  }
+}