@mastra/chroma 0.1.6-alpha.0 → 0.1.6-alpha.3

package/src/vector/index.ts

@@ -1,9 +1,25 @@
-import { type Filter } from '@mastra/core/filter';
-import { MastraVector, type QueryResult, type IndexStats } from '@mastra/core/vector';
+import { MastraVector } from '@mastra/core/vector';
+import type {
+  QueryResult,
+  IndexStats,
+  CreateIndexParams,
+  UpsertVectorParams,
+  QueryVectorParams,
+  ParamsToArgs,
+} from '@mastra/core/vector';
+import type { VectorFilter } from '@mastra/core/vector/filter';
 import { ChromaClient } from 'chromadb';
 
 import { ChromaFilterTranslator } from './filter';
 
+interface ChromaUpsertVectorParams extends UpsertVectorParams {
+  documents?: string[];
+}
+
+interface ChromaQueryVectorParams extends QueryVectorParams {
+  documentFilter?: VectorFilter;
+}
+
 export class ChromaVector extends MastraVector {
   private client: ChromaClient;
   private collections: Map<string, any>;
@@ -26,11 +42,11 @@ export class ChromaVector extends MastraVector {
     this.collections = new Map();
   }
 
-  private async getCollection(indexName: string, throwIfNotExists: boolean = true) {
+  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) {
+    } catch {
       if (throwIfNotExists) {
         throw new Error(`Index ${indexName} does not exist`);
       }
@@ -49,12 +65,11 @@ export class ChromaVector extends MastraVector {
     }
   }
 
-  async upsert(
-    indexName: string,
-    vectors: number[][],
-    metadata?: Record<string, any>[],
-    ids?: string[],
-  ): Promise<string[]> {
+  async upsert(...args: ParamsToArgs<ChromaUpsertVectorParams>): Promise<string[]> {
+    const params = this.normalizeArgs<ChromaUpsertVectorParams>('upsert', args, ['documents']);
+
+    const { indexName, vectors, metadata, ids, documents } = params;
+
     const collection = await this.getCollection(indexName);
 
     // Get index stats to check dimension
@@ -73,50 +88,60 @@ export class ChromaVector extends MastraVector {
       ids: generatedIds,
       embeddings: vectors,
       metadatas: normalizedMetadata,
+      documents: documents,
     });
 
     return generatedIds;
   }
 
-  async createIndex(
-    indexName: string,
-    dimension: number,
-    metric: 'cosine' | 'euclidean' | 'dotproduct' = 'cosine',
-  ): Promise<void> {
+  private HnswSpaceMap = {
+    cosine: 'cosine',
+    euclidean: 'l2',
+    dotproduct: 'ip',
+    l2: 'euclidean',
+    ip: 'dotproduct',
+  };
+
+  async createIndex(...args: ParamsToArgs<CreateIndexParams>): Promise<void> {
+    const params = this.normalizeArgs<CreateIndexParams>('createIndex', args);
+
+    const { indexName, dimension, metric = 'cosine' } = params;
+
     if (!Number.isInteger(dimension) || dimension <= 0) {
       throw new Error('Dimension must be a positive integer');
     }
+    const hnswSpace = this.HnswSpaceMap[metric];
+    if (!['cosine', 'l2', 'ip'].includes(hnswSpace)) {
+      throw new Error(`Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`);
+    }
     await this.client.createCollection({
       name: indexName,
       metadata: {
         dimension,
-        metric,
+        'hnsw:space': this.HnswSpaceMap[metric],
       },
     });
   }
 
-  transformFilter(filter?: Filter) {
-    const chromaFilter = new ChromaFilterTranslator();
-    const translatedFilter = chromaFilter.translate(filter);
-    return translatedFilter;
+  transformFilter(filter?: VectorFilter) {
+    const translator = new ChromaFilterTranslator();
+    return translator.translate(filter);
   }
-  async query(
-    indexName: string,
-    queryVector: number[],
-    topK: number = 10,
-    filter?: Filter,
-    includeVector: boolean = false,
-  ): Promise<QueryResult[]> {
+  async query(...args: ParamsToArgs<ChromaQueryVectorParams>): Promise<QueryResult[]> {
+    const params = this.normalizeArgs<ChromaQueryVectorParams>('query', args, ['documentFilter']);
+
+    const { indexName, queryVector, topK = 10, filter, includeVector = false, documentFilter } = params;
+
     const collection = await this.getCollection(indexName, true);
 
     const defaultInclude = ['documents', 'metadatas', 'distances'];
 
     const translatedFilter = this.transformFilter(filter);
-
     const results = await collection.query({
       queryEmbeddings: [queryVector],
       nResults: topK,
       where: translatedFilter,
+      whereDocument: documentFilter,
       include: includeVector ? [...defaultInclude, 'embeddings'] : defaultInclude,
     });
 
@@ -125,6 +150,7 @@ export class ChromaVector extends MastraVector {
       id,
       score: results.distances?.[0]?.[index] || 0,
       metadata: results.metadatas?.[0]?.[index] || {},
+      document: results.documents?.[0]?.[index],
       ...(includeVector && { vector: results.embeddings?.[0]?.[index] || [] }),
     }));
   }
@@ -139,10 +165,12 @@ export class ChromaVector extends MastraVector {
     const count = await collection.count();
     const metadata = collection.metadata;
 
+    const hnswSpace = metadata?.['hnsw:space'] as 'cosine' | 'l2' | 'ip';
+
     return {
       dimension: metadata?.dimension || 0,
       count,
-      metric: metadata?.metric as 'cosine' | 'euclidean' | 'dotproduct',
+      metric: this.HnswSpaceMap[hnswSpace] as 'cosine' | 'euclidean' | 'dotproduct',
     };
   }