@mastra/chroma 0.11.3 → 0.11.4

package/src/vector/filter.ts

@@ -21,17 +21,6 @@ export type ChromaVectorFilter = VectorFilter<
   ChromaBlacklisted
 >;
 
-type ChromaDocumentOperatorValueMap = ChromaOperatorValueMap;
-
-type ChromaDocumentBlacklisted = Exclude<ChromaBlacklisted, '$contains'>;
-
-export type ChromaVectorDocumentFilter = VectorFilter<
-  keyof ChromaDocumentOperatorValueMap,
-  ChromaDocumentOperatorValueMap,
-  ChromaLogicalOperatorValueMap,
-  ChromaDocumentBlacklisted
->;
-
 /**
  * Translator for Chroma filter queries.
  * Maintains MongoDB-compatible syntax while ensuring proper validation
@@ -59,7 +48,7 @@ export class ChromaFilterTranslator extends BaseFilterTranslator<ChromaVectorFil
   private translateNode(node: ChromaVectorFilter, currentPath: string = ''): any {
     // Handle primitive values and arrays
     if (this.isRegex(node)) {
-      throw new Error('Regex is not supported in Chroma');
+      throw new Error('Regex is supported in Chroma via the `documentFilter` argument');
     }
     if (this.isPrimitive(node)) return this.normalizeComparisonValue(node);
     if (Array.isArray(node)) return { $in: this.normalizeArrayValues(node) };
@@ -70,6 +59,9 @@ export class ChromaFilterTranslator extends BaseFilterTranslator<ChromaVectorFil
     if (entries.length === 1 && firstEntry && this.isOperator(firstEntry[0])) {
       const [operator, value] = firstEntry;
       const translated = this.translateOperator(operator, value);
+      if (this.isLogicalOperator(operator) && Array.isArray(translated) && translated.length === 1) {
+        return translated[0];
+      }
       return this.isLogicalOperator(operator) ? { [operator]: translated } : translated;
     }
 

package/src/vector/index.test.ts

@@ -1,12 +1,11 @@
+import { createVectorTestSuite } from '@internal/storage-test-utils';
 import type { QueryResult, IndexStats } from '@mastra/core/vector';
 import { describe, expect, beforeEach, afterEach, it, beforeAll, afterAll, vi } from 'vitest';
 
 import { ChromaVector } from './';
 
 describe('ChromaVector Integration Tests', () => {
-  let vectorDB = new ChromaVector({
-    path: 'http://localhost:8000',
-  });
+  let vectorDB = new ChromaVector();
 
   const testIndexName = 'test-index';
   const testIndexName2 = 'test-index-2';
@@ -122,7 +121,7 @@ describe('ChromaVector Integration Tests', () => {
       const ids = await vectorDB.upsert({ indexName: testIndexName, vectors: testVectors });
       expect(ids).toHaveLength(3);
 
-      const idToBeUpdated = ids[0];
+      const idToBeUpdated = ids[0] as string;
       const newVector = [1, 2, 3];
       const newMetaData = {
         test: 'updates',
@@ -150,7 +149,7 @@ describe('ChromaVector Integration Tests', () => {
       const ids = await vectorDB.upsert({ indexName: testIndexName, vectors: testVectors });
       expect(ids).toHaveLength(3);
 
-      const idToBeUpdated = ids[0];
+      const idToBeUpdated = ids[0] as string;
       const newMetaData = {
         test: 'updates',
       };
@@ -163,7 +162,7 @@ describe('ChromaVector Integration Tests', () => {
 
       const results: QueryResult[] = await vectorDB.query({
         indexName: testIndexName,
-        queryVector: testVectors[0],
+        queryVector: testVectors[0] as number[],
         topK: 2,
         includeVector: true,
       });
@@ -176,7 +175,7 @@ describe('ChromaVector Integration Tests', () => {
       const ids = await vectorDB.upsert({ indexName: testIndexName, vectors: testVectors });
       expect(ids).toHaveLength(3);
 
-      const idToBeUpdated = ids[0];
+      const idToBeUpdated = ids[0] as string;
       const newVector = [1, 2, 3];
 
       const update = {
@@ -204,7 +203,7 @@ describe('ChromaVector Integration Tests', () => {
     it('should delete the vector by id', async () => {
       const ids = await vectorDB.upsert({ indexName: testIndexName, vectors: testVectors });
       expect(ids).toHaveLength(3);
-      const idToBeDeleted = ids[0];
+      const idToBeDeleted = ids[0] as string;
 
       await vectorDB.deleteVector({ indexName: testIndexName, id: idToBeDeleted });
 
@@ -1285,7 +1284,8 @@ describe('ChromaVector Integration Tests', () => {
         const results = await vectorDB.query({ indexName: testIndexName3, queryVector: [1.0, 0.0, 0.0], topK: 3 });
         expect(results).toHaveLength(3);
         // Verify documents are returned
-        expect(results[0].document).toBe(testDocuments[0]);
+
+        expect(results[0]!.document).toBe(testDocuments[0]);
       });
 
       it('should filter documents using $contains', async () => {
@@ -1317,8 +1317,19 @@ describe('ChromaVector Integration Tests', () => {
           documentFilter: { $contains: 'fox' },
         });
         expect(results).toHaveLength(1);
-        expect(results[0].metadata?.source).toBe('pangram1');
-        expect(results[0].document).toContain('fox');
+        expect(results[0]!.metadata?.source).toBe('pangram1');
+        expect(results[0]!.document).toContain('fox');
+      });
+
+      it('should get records with metadata and document filters', async () => {
+        const results = await vectorDB.get({
+          indexName: testIndexName3,
+          filter: { source: 'pangram1' },
+          documentFilter: { $contains: 'fox' },
+        });
+        expect(results).toHaveLength(1);
+        expect(results[0]!.metadata?.source).toBe('pangram1');
+        expect(results[0]!.document).toContain('fox');
       });
     });
 
@@ -1331,8 +1342,8 @@ describe('ChromaVector Integration Tests', () => {
           documentFilter: { $and: [{ $contains: 'quick' }, { $not_contains: 'fox' }] },
         });
         expect(results).toHaveLength(1);
-        expect(results[0].document).toContain('quick');
-        expect(results[0].document).not.toContain('fox');
+        expect(results[0]!.document).toContain('quick');
+        expect(results[0]!.document).not.toContain('fox');
       });
 
       it('should handle $or conditions', async () => {
@@ -1343,8 +1354,8 @@ describe('ChromaVector Integration Tests', () => {
           documentFilter: { $or: [{ $contains: 'fox' }, { $contains: 'zebras' }] },
         });
         expect(results).toHaveLength(2);
-        expect(results[0].document).toContain('fox');
-        expect(results[1].document).toContain('zebras');
+        expect(results[0]!.document).toContain('fox');
+        expect(results[1]!.document).toContain('zebras');
       });
     });
 
@@ -1395,7 +1406,7 @@ describe('ChromaVector Integration Tests', () => {
           documentFilter: { $contains: 'quick brown' }, // Test multi-word match
         });
         expect(results.length).toBe(1);
-        expect(results[0].document).toContain('quick brown');
+        expect(results[0]!.document).toContain('quick brown');
       });
 
       it('should handle deeply nested logical operators', async () => {
@@ -1442,6 +1453,7 @@ describe('ChromaVector Integration Tests', () => {
           vectorDB.query({
             indexName: testIndexName3,
             queryVector: [1, 0, 0],
+            // @ts-ignore
             documentFilter: {},
           }),
         ).rejects.toThrow();
@@ -1548,3 +1560,113 @@ describe('ChromaVector Integration Tests', () => {
     }, 30000);
   });
 });
+
+// Metadata filtering tests for Memory system
+describe('Chroma Metadata Filtering', () => {
+  const chromaVector = new ChromaVector();
+
+  createVectorTestSuite({
+    vector: chromaVector,
+    createIndex: async (indexName: string) => {
+      // Using dimension 4 as required by the metadata filtering test vectors
+      await chromaVector.createIndex({ indexName, dimension: 4 });
+    },
+    deleteIndex: async (indexName: string) => {
+      await chromaVector.deleteIndex({ indexName });
+    },
+    waitForIndexing: async () => {
+      // Chroma may need a short wait for indexing
+      await new Promise(resolve => setTimeout(resolve, 2000));
+    },
+  });
+});
+
+// ChromaCloudVector fork functionality tests (requires CHROMA_API_KEY)
+describe.skipIf(!process.env.CHROMA_API_KEY)('ChromaCloudVector Fork Tests', () => {
+  let cloudVector: ChromaVector;
+  const testIndexName = 'fork-test-index';
+  const forkedIndexName = 'forked-test-index';
+  const dimension = 3;
+
+  beforeEach(async () => {
+    cloudVector = new ChromaVector({
+      apiKey: process.env.CHROMA_API_KEY,
+    });
+
+    // Clean up any existing test indexes
+    try {
+      await cloudVector.deleteIndex({ indexName: testIndexName });
+    } catch {
+      // Ignore errors if index doesn't exist
+    }
+    try {
+      await cloudVector.deleteIndex({ indexName: forkedIndexName });
+    } catch {
+      // Ignore errors if index doesn't exist
+    }
+  });
+
+  afterEach(async () => {
+    // Clean up test indexes
+    try {
+      await cloudVector.deleteIndex({ indexName: testIndexName });
+    } catch {
+      // Ignore cleanup errors
+    }
+    try {
+      await cloudVector.deleteIndex({ indexName: forkedIndexName });
+    } catch {
+      // Ignore cleanup errors
+    }
+  });
+
+  it('should fork an index successfully', async () => {
+    // Create initial index with some data
+    await cloudVector.createIndex({ indexName: testIndexName, dimension });
+
+    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'];
+
+    await cloudVector.upsert({
+      indexName: testIndexName,
+      vectors: testVectors,
+      ids: testIds,
+      metadata: testMetadata,
+    });
+
+    // Fork the index
+    await cloudVector.forkIndex({
+      indexName: testIndexName,
+      newIndexName: forkedIndexName,
+    });
+
+    // Verify both indexes exist and have the same data
+    let originalStats = await cloudVector.describeIndex({ indexName: testIndexName });
+    let forkedStats = await cloudVector.describeIndex({ indexName: forkedIndexName });
+
+    expect(originalStats.count).toBe(3);
+    expect(forkedStats.count).toBe(3);
+
+    await cloudVector.deleteVector({ indexName: forkedIndexName, id: 'vec1' });
+
+    originalStats = await cloudVector.describeIndex({ indexName: testIndexName });
+    forkedStats = await cloudVector.describeIndex({ indexName: forkedIndexName });
+
+    expect(originalStats.count).toBe(3);
+    expect(forkedStats.count).toBe(2);
+  });
+
+  it('should throw error when forking non-existent index', async () => {
+    await expect(
+      cloudVector.forkIndex({
+        indexName: 'non-existent-index',
+        newIndexName: forkedIndexName,
+      }),
+    ).rejects.toThrow();
+  });
+});

package/src/vector/index.ts

@@ -11,9 +11,9 @@ import type {
   DeleteVectorParams,
   UpdateVectorParams,
 } from '@mastra/core/vector';
-import { ChromaClient } from 'chromadb';
-import type { UpdateRecordsParams, Collection } from 'chromadb';
-import type { ChromaVectorDocumentFilter, ChromaVectorFilter } from './filter';
+import { ChromaClient, CloudClient } from 'chromadb';
+import type { ChromaClientArgs, RecordSet, Where, WhereDocument, Collection, Metadata } from 'chromadb';
+import type { ChromaVectorFilter } from './filter';
 import { ChromaFilterTranslator } from './filter';
 
 interface ChromaUpsertVectorParams extends UpsertVectorParams {
@@ -21,42 +21,68 @@ interface ChromaUpsertVectorParams extends UpsertVectorParams {
 }
 
 interface ChromaQueryVectorParams extends QueryVectorParams<ChromaVectorFilter> {
-  documentFilter?: ChromaVectorDocumentFilter;
+  documentFilter?: WhereDocument | null;
 }
 
+interface ChromaGetRecordsParams {
+  indexName: string;
+  ids?: string[];
+  filter?: ChromaVectorFilter;
+  documentFilter?: WhereDocument | null;
+  includeVector?: boolean;
+  limit?: number;
+  offset?: number;
+}
+
+type MastraMetadata = {
+  dimension?: number;
+};
+
+type ChromaVectorArgs = ChromaClientArgs & { apiKey?: string };
+
+const spaceMappings = {
+  cosine: 'cosine',
+  euclidean: 'l2',
+  dotproduct: 'ip',
+  l2: 'euclidean',
+  ip: 'dotproduct',
+};
+
 export class ChromaVector extends MastraVector<ChromaVectorFilter> {
   private client: ChromaClient;
-  private collections: Map<string, any>;
-
-  constructor({
-    path,
-    auth,
-  }: {
-    path: string;
-    auth?: {
-      provider: string;
-      credentials: string;
-    };
-  }) {
+  private collections: Map<string, Collection>;
+
+  constructor(chromaClientArgs?: ChromaVectorArgs) {
     super();
-    this.client = new ChromaClient({
-      path,
-      auth,
-    });
+    if (chromaClientArgs?.apiKey) {
+      this.client = new CloudClient({
+        apiKey: chromaClientArgs.apiKey,
+        tenant: chromaClientArgs.tenant,
+        database: chromaClientArgs.database,
+      });
+    } else {
+      this.client = new ChromaClient(chromaClientArgs);
+    }
     this.collections = new Map();
   }
 
-  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 {
-      if (throwIfNotExists) {
-        throw new Error(`Index ${indexName} does not exist`);
+  async getCollection({ indexName, forceUpdate = false }: { indexName: string; forceUpdate?: boolean }) {
+    let collection = this.collections.get(indexName);
+    if (forceUpdate || !collection) {
+      try {
+        collection = await this.client.getCollection({ name: indexName });
+        this.collections.set(indexName, collection);
+        return collection;
+      } catch {
+        throw new MastraError({
+          id: 'CHROMA_COLLECTION_GET_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        });
       }
-      return null;
     }
-    return this.collections.get(indexName);
+    return collection;
   }
 
   private validateVectorDimensions(vectors: number[][], dimension: number): void {
@@ -71,17 +97,16 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
 
   async upsert({ indexName, vectors, metadata, ids, documents }: ChromaUpsertVectorParams): Promise<string[]> {
     try {
-      const collection = await this.getCollection(indexName);
+      const collection = await this.getCollection({ indexName });
 
       const stats = await this.describeIndex({ indexName });
       this.validateVectorDimensions(vectors, stats.dimension);
       const generatedIds = ids || vectors.map(() => crypto.randomUUID());
-      const normalizedMetadata = metadata || vectors.map(() => ({}));
 
       await collection.upsert({
         ids: generatedIds,
         embeddings: vectors,
-        metadatas: normalizedMetadata,
+        metadatas: metadata,
         documents: documents,
       });
 
@@ -100,14 +125,6 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
     }
   }
 
-  private HnswSpaceMap = {
-    cosine: 'cosine',
-    euclidean: 'l2',
-    dotproduct: 'ip',
-    l2: 'euclidean',
-    ip: 'dotproduct',
-  };
-
   async createIndex({ indexName, dimension, metric = 'cosine' }: CreateIndexParams): Promise<void> {
     if (!Number.isInteger(dimension) || dimension <= 0) {
       throw new MastraError({
@@ -118,7 +135,9 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
         details: { dimension },
       });
     }
-    const hnswSpace = this.HnswSpaceMap[metric];
+
+    const hnswSpace = spaceMappings[metric] as 'cosine' | 'l2' | 'ip' | undefined;
+
     if (!hnswSpace || !['cosine', 'l2', 'ip'].includes(hnswSpace)) {
       throw new MastraError({
         id: 'CHROMA_VECTOR_CREATE_INDEX_INVALID_METRIC',
@@ -128,14 +147,15 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
         details: { metric },
       });
     }
+
     try {
-      await this.client.createCollection({
+      const collection = await this.client.createCollection({
         name: indexName,
-        metadata: {
-          dimension,
-          'hnsw:space': hnswSpace,
-        },
+        metadata: { dimension },
+        configuration: { hnsw: { space: hnswSpace } },
+        embeddingFunction: null,
       });
+      this.collections.set(indexName, collection);
     } catch (error: any) {
       // Check for 'already exists' error
       const message = error?.message || error?.toString();
@@ -158,9 +178,11 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
 
   transformFilter(filter?: ChromaVectorFilter) {
     const translator = new ChromaFilterTranslator();
-    return translator.translate(filter);
+    const translatedFilter = translator.translate(filter);
+    return translatedFilter ? (translatedFilter as Where) : undefined;
   }
-  async query({
+
+  async query<T extends Metadata = Metadata>({
     indexName,
     queryVector,
     topK = 10,
@@ -169,16 +191,16 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
     documentFilter,
   }: ChromaQueryVectorParams): Promise<QueryResult[]> {
     try {
-      const collection = await this.getCollection(indexName, true);
+      const collection = await this.getCollection({ indexName });
 
-      const defaultInclude = ['documents', 'metadatas', 'distances'];
+      const defaultInclude: ['documents', 'metadatas', 'distances'] = ['documents', 'metadatas', 'distances'];
 
       const translatedFilter = this.transformFilter(filter);
-      const results = await collection.query({
+      const results = await collection.query<T>({
         queryEmbeddings: [queryVector],
         nResults: topK,
-        where: translatedFilter,
-        whereDocument: documentFilter,
+        where: translatedFilter ?? undefined,
+        whereDocument: documentFilter ?? undefined,
         include: includeVector ? [...defaultInclude, 'embeddings'] : defaultInclude,
       });
 
@@ -186,7 +208,7 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
         id,
         score: results.distances?.[0]?.[index] || 0,
         metadata: results.metadatas?.[0]?.[index] || {},
-        document: results.documents?.[0]?.[index],
+        document: results.documents?.[0]?.[index] ?? undefined,
         ...(includeVector && { vector: results.embeddings?.[0]?.[index] || [] }),
       }));
     } catch (error: any) {
@@ -203,10 +225,48 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
     }
   }
 
+  async get<T extends Metadata = Metadata>({
+    indexName,
+    ids,
+    filter,
+    includeVector = false,
+    documentFilter,
+    offset,
+    limit,
+  }: ChromaGetRecordsParams) {
+    try {
+      const collection = await this.getCollection({ indexName });
+
+      const defaultInclude: ['documents', 'metadatas'] = ['documents', 'metadatas'];
+      const translatedFilter = this.transformFilter(filter);
+
+      const result = await collection.get<T>({
+        ids,
+        where: translatedFilter ?? undefined,
+        whereDocument: documentFilter ?? undefined,
+        offset,
+        limit,
+        include: includeVector ? [...defaultInclude, 'embeddings'] : defaultInclude,
+      });
+      return result.rows();
+    } catch (error: any) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_GET_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
+  }
+
   async listIndexes(): Promise<string[]> {
     try {
       const collections = await this.client.listCollections();
-      return collections.map(collection => collection);
+      return collections.map(collection => collection.name);
     } catch (error: any) {
       throw new MastraError(
         {
@@ -227,16 +287,15 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
    */
   async describeIndex({ indexName }: DescribeIndexParams): Promise<IndexStats> {
     try {
-      const collection = await this.getCollection(indexName);
+      const collection = await this.getCollection({ indexName });
       const count = await collection.count();
-      const metadata = collection.metadata;
-
-      const hnswSpace = metadata?.['hnsw:space'] as 'cosine' | 'l2' | 'ip';
+      const metadata = collection.metadata as MastraMetadata | undefined;
+      const space = collection.configuration.hnsw?.space || collection.configuration.spann?.space || undefined;
 
       return {
         dimension: metadata?.dimension || 0,
         count,
-        metric: this.HnswSpaceMap[hnswSpace] as 'cosine' | 'euclidean' | 'dotproduct',
+        metric: space ? (spaceMappings[space] as 'cosine' | 'euclidean' | 'dotproduct') : undefined,
       };
     } catch (error: any) {
       if (error instanceof MastraError) throw error;
@@ -269,6 +328,25 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
     }
   }
 
+  async forkIndex({ indexName, newIndexName }: { indexName: string; newIndexName: string }): Promise<void> {
+    try {
+      const collection = await this.getCollection({ indexName, forceUpdate: true });
+      const forkedCollection = await collection.fork({ name: newIndexName });
+      this.collections.set(newIndexName, forkedCollection);
+    } catch (error: any) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_INDEX_FORK_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
+  }
+
   /**
    * Updates a vector by its ID with the provided vector and/or metadata.
    * @param indexName - The name of the index containing the vector.
@@ -291,21 +369,21 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
     }
 
     try {
-      const collection: Collection = await this.getCollection(indexName, true);
+      const collection: Collection = await this.getCollection({ indexName });
 
-      const updateOptions: UpdateRecordsParams = { ids: [id] };
+      const updateRecordSet: RecordSet = { ids: [id] };
 
       if (update?.vector) {
         const stats = await this.describeIndex({ indexName });
         this.validateVectorDimensions([update.vector], stats.dimension);
-        updateOptions.embeddings = [update.vector];
+        updateRecordSet.embeddings = [update.vector];
       }
 
       if (update?.metadata) {
-        updateOptions.metadatas = [update.metadata];
+        updateRecordSet.metadatas = [update.metadata];
       }
 
-      return await collection.update(updateOptions);
+      return await collection.update(updateRecordSet);
     } catch (error: any) {
       if (error instanceof MastraError) throw error;
       throw new MastraError(
@@ -322,7 +400,7 @@ export class ChromaVector extends MastraVector<ChromaVectorFilter> {
 
   async deleteVector({ indexName, id }: DeleteVectorParams): Promise<void> {
     try {
-      const collection: Collection = await this.getCollection(indexName, true);
+      const collection: Collection = await this.getCollection({ indexName });
       await collection.delete({ ids: [id] });
     } catch (error: any) {
       if (error instanceof MastraError) throw error;