@mastra/rag 0.0.2-alpha.13 → 0.0.2-alpha.15

package/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @mastra/rag
 
+## 0.0.2-alpha.15
+
+### Patch Changes
+
+- Updated dependencies [c2dd6b5]
+  - @mastra/core@0.1.27-alpha.31
+
+## 0.0.2-alpha.14
+
+### Patch Changes
+
+- 1c3232a: ChromaDB
+
 ## 0.0.2-alpha.13
 
 ### Patch Changes

package/dist/chroma/index.d.ts

@@ -0,0 +1,20 @@
+import { MastraVector, QueryResult, IndexStats } from '@mastra/core';
+export declare class ChromaVector extends MastraVector {
+    private client;
+    private collections;
+    constructor({ path, auth, }: {
+        path: string;
+        auth?: {
+            provider: string;
+            credentials: string;
+        };
+    });
+    private getCollection;
+    private validateVectorDimensions;
+    upsert(indexName: string, vectors: number[][], metadata?: Record<string, any>[], ids?: string[]): Promise<string[]>;
+    createIndex(indexName: string, dimension: number, metric?: 'cosine' | 'euclidean' | 'dotproduct'): Promise<void>;
+    query(indexName: string, queryVector: number[], topK?: number, filter?: Record<string, any>): Promise<QueryResult[]>;
+    listIndexes(): Promise<string[]>;
+    describeIndex(indexName: string): Promise<IndexStats>;
+    deleteIndex(indexName: string): Promise<void>;
+}

package/docker-compose.yaml

@@ -14,5 +14,9 @@ services:
     image: qdrant/qdrant
     ports:
       - 6333:6333
+  chroma:
+    image: chromadb/chroma
+    ports:
+      - 8000:8000
 volumes:
   pgdata:

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/rag",
-  "version": "0.0.2-alpha.13",
+  "version": "0.0.2-alpha.15",
   "description": "",
   "type": "module",
   "main": "dist/index.js",
@@ -32,13 +32,15 @@
     "@pinecone-database/pinecone": "^3.0.3",
     "@qdrant/js-client-rest": "^1.12.0",
     "@upstash/vector": "^1.1.7",
+    "chromadb": "^1.9.4",
+    "chromadb-default-embed": "^2.13.2",
     "date-fns": "^4.1.0",
     "dotenv": "^16.3.1",
     "jsdom": "^25.0.1",
     "llamaindex": "^0.8.15",
     "pg": "^8.13.1",
     "postgres": "^3.4.5",
-    "@mastra/core": "0.1.27-alpha.30"
+    "@mastra/core": "0.1.27-alpha.31"
   },
   "devDependencies": {
     "@babel/preset-env": "^7.26.0",

package/src/chroma/index.test.ts

@@ -0,0 +1,154 @@
+import { QueryResult, IndexStats } from '@mastra/core';
+
+import { ChromaVector } from './';
+
+describe('ChromaVector Integration Tests', () => {
+  let vectorDB: ChromaVector;
+  const testIndexName = 'test-index';
+  const dimension = 3;
+
+  beforeEach(async () => {
+    vectorDB = new ChromaVector({
+      path: 'http://localhost:8000',
+    });
+    // Clean up any existing test index
+    try {
+      await vectorDB.deleteIndex(testIndexName);
+    } catch (error) {
+      // Ignore errors if index doesn't exist
+    }
+    await vectorDB.createIndex(testIndexName, dimension);
+  }, 5000);
+
+  afterEach(async () => {
+    // Cleanup after tests
+    try {
+      await vectorDB.deleteIndex(testIndexName);
+    } catch (error) {
+      // Ignore cleanup errors
+    }
+  }, 5000);
+
+  describe('Index Management', () => {
+    test('should create and list indexes', async () => {
+      const indexes = await vectorDB.listIndexes();
+      expect(indexes).toContain(testIndexName);
+    });
+
+    test('should describe index correctly', async () => {
+      const stats: IndexStats = await vectorDB.describeIndex(testIndexName);
+      expect(stats.dimension).toBe(dimension);
+      expect(stats.count).toBe(0);
+      expect(stats.metric).toBe('cosine');
+    });
+
+    test('should delete index', async () => {
+      await vectorDB.deleteIndex(testIndexName);
+      const indexes = await vectorDB.listIndexes();
+      expect(indexes).not.toContain(testIndexName);
+    });
+
+    test('should create index with different metrics', async () => {
+      const metricsToTest: Array<'cosine' | 'euclidean' | 'dotproduct'> = ['euclidean', 'dotproduct'];
+
+      for (const metric of metricsToTest) {
+        const testIndex = `test-index-${metric}`;
+        await vectorDB.createIndex(testIndex, dimension, metric);
+
+        const stats = await vectorDB.describeIndex(testIndex);
+        expect(stats.metric).toBe(metric);
+
+        await vectorDB.deleteIndex(testIndex);
+      }
+    });
+  });
+
+  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' }];
+    const testIds = ['vec1', 'vec2', 'vec3'];
+
+    test('should upsert vectors with generated ids', async () => {
+      const ids = await vectorDB.upsert(testIndexName, testVectors);
+      expect(ids).toHaveLength(testVectors.length);
+      ids.forEach(id => expect(typeof id).toBe('string'));
+
+      const stats = await vectorDB.describeIndex(testIndexName);
+      expect(stats.count).toBe(testVectors.length);
+    });
+
+    test('should upsert vectors with provided ids and metadata', async () => {
+      await vectorDB.upsert(testIndexName, testVectors, testMetadata, testIds);
+
+      const stats = await vectorDB.describeIndex(testIndexName);
+      expect(stats.count).toBe(testVectors.length);
+
+      // Query each vector to verify metadata
+      for (let i = 0; i < testVectors.length; i++) {
+        const results = await vectorDB.query(testIndexName, testVectors?.[i]!, 1);
+        expect(results?.[0]?.id).toBe(testIds[i]);
+        expect(results?.[0]?.metadata).toEqual(testMetadata[i]);
+      }
+    });
+
+    test('should perform vector search with topK', async () => {
+      await vectorDB.upsert(testIndexName, testVectors, testMetadata, testIds);
+
+      const queryVector = [1.0, 0.1, 0.1];
+      const topK = 2;
+
+      const results: QueryResult[] = await vectorDB.query(testIndexName, queryVector, topK);
+
+      expect(results).toHaveLength(topK);
+      expect(results?.[0]?.id).toBe(testIds[0]); // Should match x-axis vector most closely
+    });
+
+    test('should filter query results', async () => {
+      await vectorDB.upsert(testIndexName, testVectors, testMetadata, testIds);
+
+      const queryVector = [1.0, 1.0, 1.0];
+      const filter = { label: 'x-axis' };
+
+      const results = await vectorDB.query(testIndexName, queryVector, 3, filter);
+
+      expect(results).toHaveLength(1);
+      expect(results?.[0]?.metadata?.label).toBe('x-axis');
+    });
+
+    test('should update existing vectors', async () => {
+      // Initial upsert
+      await vectorDB.upsert(testIndexName, testVectors, testMetadata, testIds);
+
+      // Update first vector
+      const updatedVector = [[0.5, 0.5, 0.0]];
+      const updatedMetadata = [{ label: 'updated-x-axis' }];
+      await vectorDB.upsert(testIndexName, updatedVector, updatedMetadata, [testIds?.[0]!]);
+
+      // Verify update
+      const results = await vectorDB.query(testIndexName, updatedVector?.[0]!, 1);
+      expect(results?.[0]?.id).toBe(testIds[0]);
+      expect(results?.[0]?.metadata).toEqual(updatedMetadata[0]);
+    });
+  });
+
+  describe('Error Handling', () => {
+    test('should handle non-existent index queries', async () => {
+      await expect(vectorDB.query('non-existent-index-yu', [1, 2, 3])).rejects.toThrow();
+    });
+
+    test('should handle invalid dimension vectors', async () => {
+      const invalidVector = [1, 2, 3, 4]; // 4D vector for 3D index
+      await expect(vectorDB.upsert(testIndexName, [invalidVector])).rejects.toThrow();
+    });
+
+    test('should handle mismatched metadata and vectors length', async () => {
+      const vectors = [[1, 2, 3]];
+      const metadata = [{}, {}]; // More metadata than vectors
+      await expect(vectorDB.upsert(testIndexName, vectors, metadata)).rejects.toThrow();
+    });
+  });
+});

package/src/chroma/index.ts

@@ -0,0 +1,135 @@
+import { MastraVector, QueryResult, IndexStats } from '@mastra/core';
+import { ChromaClient } from 'chromadb';
+
+export class ChromaVector extends MastraVector {
+  private client: ChromaClient;
+  private collections: Map<string, any>;
+
+  constructor({
+    path,
+    auth,
+  }: {
+    path: string;
+    auth?: {
+      provider: string;
+      credentials: string;
+    };
+  }) {
+    super();
+    this.client = new ChromaClient({
+      path,
+      auth,
+    });
+    this.collections = new Map();
+  }
+
+  private async getCollection(indexName: string, throwIfNotExists: boolean = true) {
+    try {
+      const collection = await this.client.getCollection({ name: indexName, embeddingFunction: undefined as any });
+      this.collections.set(indexName, collection);
+    } catch (error) {
+      if (throwIfNotExists) {
+        throw new Error(`Index ${indexName} does not exist`);
+      }
+      return null;
+    }
+    return this.collections.get(indexName);
+  }
+
+  private validateVectorDimensions(vectors: number[][], dimension: number): void {
+    for (let i = 0; i < vectors.length; i++) {
+      if (vectors?.[i]?.length !== dimension) {
+        throw new Error(
+          `Vector at index ${i} has invalid dimension ${vectors?.[i]?.length}. Expected ${dimension} dimensions.`,
+        );
+      }
+    }
+  }
+
+  async upsert(
+    indexName: string,
+    vectors: number[][],
+    metadata?: Record<string, any>[],
+    ids?: string[],
+  ): Promise<string[]> {
+    const collection = await this.getCollection(indexName);
+
+    // Get index stats to check dimension
+    const stats = await this.describeIndex(indexName);
+
+    // Validate vector dimensions
+    this.validateVectorDimensions(vectors, stats.dimension);
+
+    // Generate IDs if not provided
+    const generatedIds = ids || vectors.map(() => crypto.randomUUID());
+
+    // Ensure metadata exists for each vector
+    const normalizedMetadata = metadata || vectors.map(() => ({}));
+
+    await collection.upsert({
+      ids: generatedIds,
+      embeddings: vectors,
+      metadatas: normalizedMetadata,
+    });
+
+    return generatedIds;
+  }
+
+  async createIndex(
+    indexName: string,
+    dimension: number,
+    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
+  ): Promise<void> {
+    await this.client.createCollection({
+      name: indexName,
+      metadata: {
+        dimension,
+        metric,
+      },
+    });
+  }
+
+  async query(
+    indexName: string,
+    queryVector: number[],
+    topK: number = 10,
+    filter?: Record<string, any>,
+  ): Promise<QueryResult[]> {
+    const collection = await this.getCollection(indexName, true);
+
+    const results = await collection.query({
+      queryEmbeddings: [queryVector],
+      nResults: topK,
+      where: filter,
+    });
+
+    // Transform ChromaDB results to QueryResult format
+    return (results.ids[0] || []).map((id: string, index: number) => ({
+      id,
+      score: results.distances?.[0]?.[index] || 0,
+      metadata: results.metadatas?.[0]?.[index] || {},
+    }));
+  }
+
+  async listIndexes(): Promise<string[]> {
+    const collections = await this.client.listCollections();
+    return collections.map(collection => collection.name);
+  }
+
+  async describeIndex(indexName: string): Promise<IndexStats> {
+    const collection = await this.getCollection(indexName);
+    const count = await collection.count();
+    const metadata = collection.metadata;
+
+    return {
+      dimension: metadata?.dimension || 0,
+      count,
+      metric: metadata?.metric as 'cosine' | 'euclidean' | 'dotproduct',
+    };
+  }
+
+  async deleteIndex(indexName: string): Promise<void> {
+    await this.client.deleteCollection({ name: indexName });
+    this.collections.delete(indexName);
+  }
+}

package/src/document/document.test.ts

@@ -10,169 +10,6 @@ Welcome to our comprehensive guide on modern web development. This resource cove
 - Beginning developers looking to establish a solid foundation
 - Intermediate developers wanting to modernize their skillset
 - Senior developers seeking a refresher on current best practices
-
-## Core Concepts
-
-### 1. Frontend Development
-Modern frontend development has evolved significantly. Here are the key areas to focus on:
-
-#### HTML5 Semantic Elements
-Using semantic HTML improves:
-- Accessibility
-- SEO performance
-- Code readability
-- Maintenance
-
-\`\`\`html
-<header>
-    <nav>
-        <ul>
-            <li><a href="#home">Home</a></li>
-            <li><a href="#about">About</a></li>
-        </ul>
-    </nav>
-</header>
-\`\`\`
-
-#### CSS Best Practices
-1. Use CSS Custom Properties
-2. Implement responsive design
-3. Follow BEM methodology
-4. Optimize performance
-
-### 2. JavaScript Fundamentals
-
-JavaScript is the backbone of web development. Here's what you need to know:
-
-\`\`\`javascript
-// Modern JS features
-const exampleFunction = async () => {
-    try {
-        const response = await fetch('https://api.example.com/data');
-        const data = await response.json();
-        return data;
-    } catch (error) {
-        console.error('Error fetching data:', error);
-    }
-};
-\`\`\`
-
-#### Key Concepts:
-- Promises and async/await
-- ES6+ features
-- TypeScript integration
-- Module systems
-
-### 3. Backend Development
-
-Backend development requires understanding:
-
-1. **Server Architecture**
-   - RESTful APIs
-   - GraphQL
-   - Microservices
-
-2. **Database Management**
-   - SQL vs NoSQL
-   - Query optimization
-   - Data modeling
-
-3. **Security Considerations**
-   - Authentication
-   - Authorization
-   - Data encryption
-
-## Tools and Technologies
-
-### Essential Developer Tools
-
-| Category | Tools |
-|----------|-------|
-| Version Control | Git, GitHub |
-| Package Managers | npm, yarn |
-| Bundlers | webpack, Vite |
-| Testing | Jest, Cypress |
-
-### Framework Selection
-
-#### Frontend Frameworks
-1. React
-   - Component-based architecture
-   - Virtual DOM
-   - Large ecosystem
-
-2. Vue
-   - Progressive framework
-   - Easy learning curve
-   - Great documentation
-
-3. Angular
-   - Full-featured framework
-   - TypeScript integration
-   - Enterprise-ready
-
-#### Backend Frameworks
-- Node.js/Express
-- Django
-- Ruby on Rails
-- Spring Boot
-
-## Best Practices
-
-### Code Quality
-- Write clean, maintainable code
-- Follow SOLID principles
-- Implement proper error handling
-- Use consistent formatting
-
-### Performance Optimization
-1. Minimize HTTP requests
-2. Optimize images
-3. Implement caching
-4. Use lazy loading
-
-### Security Measures
-- Implement HTTPS
-- Sanitize user input
-- Use secure dependencies
-- Regular security audits
-
-## Deployment and DevOps
-
-### Continuous Integration/Continuous Deployment (CI/CD)
-1. Automated testing
-2. Build automation
-3. Deployment automation
-4. Monitoring and logging
-
-### Cloud Services
-- AWS
-- Google Cloud Platform
-- Azure
-- Heroku
-
-## Conclusion
-
-Remember that web development is an ever-evolving field. Stay current with:
-- Industry trends
-- New tools and frameworks
-- Security best practices
-- Performance optimization techniques
-
-### Additional Resources
-
-* [MDN Web Docs](https://developer.mozilla.org)
-* [Web.dev](https://web.dev)
-* [CSS-Tricks](https://css-tricks.com)
-* [JavaScript.info](https://javascript.info)
-
----
-
-> "Any application that can be written in JavaScript, will eventually be written in JavaScript." - Jeff Atwood
-
----
-
-**Note**: This guide is regularly updated to reflect current web development practices and standards.
 `;
 
 describe('MastraDocument', () => {
@@ -195,19 +32,17 @@ describe('MastraDocument', () => {
 
       await doc.chunk({
         extract: {
-          summary: true,
-          title: true,
+          keywords: true,
         },
         options: {
-          chunkSize: 100,
+          chunkSize: 1500,
           chunkOverlap: 0,
           separator: `\n`,
         },
       });
 
-      expect(doc.getDocs()?.[0]?.text).toBe(`Complete Guide to Modern Web Development`);
-      expect(doc.getDocs().length).toBe(56);
-    }, 500000);
+      expect(doc.getMetadata()?.[0]).toBeTruthy();
+    }, 15000);
   });
 
   describe('chunkCharacter', () => {
@@ -672,7 +507,6 @@ describe('MastraDocument', () => {
             const c = JSON.parse(chunk);
             const retVal: Record<string, string> = {};
             Object.entries(c).forEach(([key, value]) => {
-              console.log(key, value);
               retVal[key] = JSON.parse(`"${value as string}"`);
             });