@mastra/chroma 0.0.0-vector-query-sources-20250516172905 → 0.0.0-vector-query-tool-provider-options-20250828222356

package/src/vector/index.ts

@@ -1,3 +1,4 @@
+import { MastraError, ErrorDomain, ErrorCategory } from '@mastra/core/error';
 import { MastraVector } from '@mastra/core/vector';
 import type {
   QueryResult,
@@ -5,65 +6,83 @@ import type {
   CreateIndexParams,
   UpsertVectorParams,
   QueryVectorParams,
-  ParamsToArgs,
-  QueryVectorArgs,
-  UpsertVectorArgs,
   DescribeIndexParams,
   DeleteIndexParams,
   DeleteVectorParams,
   UpdateVectorParams,
 } from '@mastra/core/vector';
-
-import type { VectorFilter } from '@mastra/core/vector/filter';
-import { ChromaClient } from 'chromadb';
-import type { UpdateRecordsParams, Collection } from 'chromadb';
+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 {
   documents?: string[];
 }
 
-type ChromaUpsertArgs = [...UpsertVectorArgs, string[]?];
+interface ChromaQueryVectorParams extends QueryVectorParams<ChromaVectorFilter> {
+  documentFilter?: WhereDocument | null;
+}
 
-interface ChromaQueryVectorParams extends QueryVectorParams {
-  documentFilter?: VectorFilter;
+interface ChromaGetRecordsParams {
+  indexName: string;
+  ids?: string[];
+  filter?: ChromaVectorFilter;
+  documentFilter?: WhereDocument | null;
+  includeVector?: boolean;
+  limit?: number;
+  offset?: number;
 }
 
-type ChromaQueryArgs = [...QueryVectorArgs, VectorFilter?];
+type MastraMetadata = {
+  dimension?: number;
+};
+
+type ChromaVectorArgs = ChromaClientArgs & { apiKey?: string };
 
-export class ChromaVector extends MastraVector {
+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 {
@@ -76,62 +95,67 @@ export class ChromaVector extends MastraVector {
     }
   }
 
-  async upsert(...args: ParamsToArgs<ChromaUpsertVectorParams> | ChromaUpsertArgs): Promise<string[]> {
-    const params = this.normalizeArgs<ChromaUpsertVectorParams, ChromaUpsertArgs>('upsert', args, ['documents']);
-
-    const { indexName, vectors, metadata, ids, documents } = params;
-
-    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());
+  async upsert({ indexName, vectors, metadata, ids, documents }: ChromaUpsertVectorParams): Promise<string[]> {
+    try {
+      const collection = await this.getCollection({ indexName });
 
-    // Ensure metadata exists for each vector
-    const normalizedMetadata = metadata || vectors.map(() => ({}));
+      const stats = await this.describeIndex({ indexName });
+      this.validateVectorDimensions(vectors, stats.dimension);
+      const generatedIds = ids || vectors.map(() => crypto.randomUUID());
 
-    await collection.upsert({
-      ids: generatedIds,
-      embeddings: vectors,
-      metadatas: normalizedMetadata,
-      documents: documents,
-    });
+      await collection.upsert({
+        ids: generatedIds,
+        embeddings: vectors,
+        metadatas: metadata,
+        documents: documents,
+      });
 
-    return generatedIds;
+      return generatedIds;
+    } catch (error: any) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_UPSERT_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
   }
 
-  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;
-
+  async createIndex({ indexName, dimension, metric = 'cosine' }: CreateIndexParams): Promise<void> {
     if (!Number.isInteger(dimension) || dimension <= 0) {
-      throw new Error('Dimension must be a positive integer');
+      throw new MastraError({
+        id: 'CHROMA_VECTOR_CREATE_INDEX_INVALID_DIMENSION',
+        text: 'Dimension must be a positive integer',
+        domain: ErrorDomain.MASTRA_VECTOR,
+        category: ErrorCategory.USER,
+        details: { dimension },
+      });
     }
-    const hnswSpace = this.HnswSpaceMap[metric];
-    if (!['cosine', 'l2', 'ip'].includes(hnswSpace)) {
-      throw new Error(`Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`);
+
+    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',
+        text: `Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`,
+        domain: ErrorDomain.MASTRA_VECTOR,
+        category: ErrorCategory.USER,
+        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();
@@ -140,101 +164,187 @@ export class ChromaVector extends MastraVector {
         await this.validateExistingIndex(indexName, dimension, metric);
         return;
       }
-      throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_CREATE_INDEX_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
     }
   }
 
-  transformFilter(filter?: VectorFilter) {
+  transformFilter(filter?: ChromaVectorFilter) {
     const translator = new ChromaFilterTranslator();
-    return translator.translate(filter);
+    const translatedFilter = translator.translate(filter);
+    return translatedFilter ? (translatedFilter as Where) : undefined;
   }
-  async query(...args: ParamsToArgs<ChromaQueryVectorParams> | ChromaQueryArgs): Promise<QueryResult[]> {
-    const params = this.normalizeArgs<ChromaQueryVectorParams, ChromaQueryArgs>('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,
-    });
-
-    // 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] || {},
-      document: results.documents?.[0]?.[index],
-      ...(includeVector && { vector: results.embeddings?.[0]?.[index] || [] }),
-    }));
+
+  async query<T extends Metadata = Metadata>({
+    indexName,
+    queryVector,
+    topK = 10,
+    filter,
+    includeVector = false,
+    documentFilter,
+  }: ChromaQueryVectorParams): Promise<QueryResult[]> {
+    try {
+      const collection = await this.getCollection({ indexName });
+
+      const defaultInclude: ['documents', 'metadatas', 'distances'] = ['documents', 'metadatas', 'distances'];
+
+      const translatedFilter = this.transformFilter(filter);
+      const results = await collection.query<T>({
+        queryEmbeddings: [queryVector],
+        nResults: topK,
+        where: translatedFilter ?? undefined,
+        whereDocument: documentFilter ?? undefined,
+        include: includeVector ? [...defaultInclude, 'embeddings'] : defaultInclude,
+      });
+
+      return (results.ids[0] || []).map((id: string, index: number) => ({
+        id,
+        score: results.distances?.[0]?.[index] || 0,
+        metadata: results.metadatas?.[0]?.[index] || {},
+        document: results.documents?.[0]?.[index] ?? undefined,
+        ...(includeVector && { vector: results.embeddings?.[0]?.[index] || [] }),
+      }));
+    } catch (error: any) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_QUERY_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
+  }
+
+  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[]> {
-    const collections = await this.client.listCollections();
-    return collections.map(collection => collection);
+    try {
+      const collections = await this.client.listCollections();
+      return collections.map(collection => collection.name);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_LIST_INDEXES_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
   }
 
   /**
    * Retrieves statistics about a vector index.
    *
-   * @param params - The parameters for describing an index
-   * @param params.indexName - The name of the index to describe
+   * @param {string} indexName - The name of the index to describe
    * @returns A promise that resolves to the index statistics including dimension, count and metric
    */
-  async describeIndex(...args: ParamsToArgs<DescribeIndexParams>): Promise<IndexStats> {
-    const params = this.normalizeArgs<DescribeIndexParams>('describeIndex', args);
-    const { indexName } = params;
-
-    const collection = await this.getCollection(indexName);
-    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: this.HnswSpaceMap[hnswSpace] as 'cosine' | 'euclidean' | 'dotproduct',
-    };
+  async describeIndex({ indexName }: DescribeIndexParams): Promise<IndexStats> {
+    try {
+      const collection = await this.getCollection({ indexName });
+      const count = await collection.count();
+      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: space ? (spaceMappings[space] as 'cosine' | 'euclidean' | 'dotproduct') : undefined,
+      };
+    } catch (error: any) {
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_DESCRIBE_INDEX_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
   }
 
-  async deleteIndex(...args: ParamsToArgs<DeleteIndexParams>): Promise<void> {
-    const params = this.normalizeArgs<DeleteIndexParams>('deleteIndex', args);
-    const { indexName } = params;
-    await this.client.deleteCollection({ name: indexName });
-    this.collections.delete(indexName);
+  async deleteIndex({ indexName }: DeleteIndexParams): Promise<void> {
+    try {
+      await this.client.deleteCollection({ name: indexName });
+      this.collections.delete(indexName);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_DELETE_INDEX_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName },
+        },
+        error,
+      );
+    }
   }
 
-  /**
-   * @deprecated Use {@link updateVector} instead. This method will be removed on May 20th, 2025.
-   *
-   * Updates a vector by its ID with the provided vector and/or metadata.
-   * @param indexName - The name of the index containing the vector.
-   * @param id - The ID of the vector to update.
-   * @param update - An object containing the vector and/or metadata to update.
-   * @param update.vector - An optional array of numbers representing the new vector.
-   * @param update.metadata - An optional record containing the new metadata.
-   * @returns A promise that resolves when the update is complete.
-   * @throws Will throw an error if no updates are provided or if the update operation fails.
-   */
-  async updateIndexById(
-    indexName: string,
-    id: string,
-    update: { vector?: number[]; metadata?: Record<string, any> },
-  ): Promise<void> {
-    this.logger.warn(
-      `Deprecation Warning: updateIndexById() is deprecated. 
-    Please use updateVector() instead. 
-    updateIndexById() will be removed on May 20th, 2025.`,
-    );
-    await this.updateVector({ indexName, id, update });
+  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,
+      );
+    }
   }
 
   /**
@@ -247,66 +357,62 @@ export class ChromaVector extends MastraVector {
    * @returns A promise that resolves when the update is complete.
    * @throws Will throw an error if no updates are provided or if the update operation fails.
    */
-  async updateVector(...args: ParamsToArgs<UpdateVectorParams>): Promise<void> {
-    const params = this.normalizeArgs<UpdateVectorParams>('updateVector', args);
-    const { indexName, id, update } = params;
-    try {
-      if (!update.vector && !update.metadata) {
-        throw new Error('No updates provided');
-      }
+  async updateVector({ indexName, id, update }: UpdateVectorParams): Promise<void> {
+    if (!update.vector && !update.metadata) {
+      throw new MastraError({
+        id: 'CHROMA_VECTOR_UPDATE_NO_PAYLOAD',
+        text: 'No updates provided for vector',
+        domain: ErrorDomain.MASTRA_VECTOR,
+        category: ErrorCategory.USER,
+        details: { indexName, id },
+      });
+    }
 
-      const collection: Collection = await this.getCollection(indexName, true);
+    try {
+      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) {
-      throw new Error(`Failed to update vector by id: ${id} for index name: ${indexName}: ${error.message}`);
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_UPDATE_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName, id },
+        },
+        error,
+      );
     }
   }
 
-  /**
-   * @deprecated Use {@link deleteVector} instead. This method will be removed on May 20th, 2025.
-   *
-   * Deletes a vector by its ID.
-   * @param indexName - The name of the index containing the vector.
-   * @param id - The ID of the vector to delete.
-   * @returns A promise that resolves when the deletion is complete.
-   * @throws Will throw an error if the deletion operation fails.
-   */
-  async deleteIndexById(indexName: string, id: string): Promise<void> {
-    this.logger.warn(
-      `Deprecation Warning: deleteIndexById() is deprecated. Please use deleteVector() instead. deleteIndexById() will be removed on May 20th.`,
-    );
-    await this.deleteVector(indexName, id);
-  }
-
-  /**
-   * Deletes a vector by its ID.
-   * @param indexName - The name of the index containing the vector.
-   * @param id - The ID of the vector to delete.
-   * @returns A promise that resolves when the deletion is complete.
-   * @throws Will throw an error if the deletion operation fails.
-   */
-  async deleteVector(...args: ParamsToArgs<DeleteVectorParams>): Promise<void> {
-    const params = this.normalizeArgs<DeleteVectorParams>('deleteVector', args);
-
-    const { indexName, id } = params;
+  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) {
-      throw new Error(`Failed to delete vector by id: ${id} for index name: ${indexName}: ${error.message}`);
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'CHROMA_VECTOR_DELETE_FAILED',
+          domain: ErrorDomain.MASTRA_VECTOR,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { indexName, id },
+        },
+        error,
+      );
     }
   }
 }

package/tsconfig.build.json

@@ -0,0 +1,9 @@
+{
+  "extends": ["./tsconfig.json", "../../tsconfig.build.json"],
+  "compilerOptions": {
+    "outDir": "./dist",
+    "rootDir": "./src"
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "**/*.test.ts", "src/**/*.mock.ts"]
+}

package/tsconfig.json

@@ -1,5 +1,5 @@
 {
   "extends": "../../tsconfig.node.json",
-  "include": ["src/**/*"],
+  "include": ["src/**/*", "tsup.config.ts"],
   "exclude": ["node_modules", "**/*.test.ts"]
 }

package/tsup.config.ts

@@ -0,0 +1,17 @@
+import { generateTypes } from '@internal/types-builder';
+import { defineConfig } from 'tsup';
+
+export default defineConfig({
+  entry: ['src/index.ts'],
+  format: ['esm', 'cjs'],
+  clean: true,
+  dts: false,
+  splitting: true,
+  treeshake: {
+    preset: 'smallest',
+  },
+  sourcemap: true,
+  onSuccess: async () => {
+    await generateTypes(process.cwd());
+  },
+});