@mastra/chroma 0.1.6-alpha.0 → 0.1.6-alpha.3
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.turbo/turbo-build.log +11 -6
- package/CHANGELOG.md +37 -0
- package/README.md +26 -16
- package/dist/_tsup-dts-rollup.d.cts +56 -0
- package/dist/_tsup-dts-rollup.d.ts +24 -11
- package/dist/index.cjs +214 -0
- package/dist/index.d.cts +1 -0
- package/dist/index.js +31 -11
- package/docker-compose.yaml +7 -0
- package/eslint.config.js +6 -0
- package/package.json +12 -5
- package/src/vector/filter.ts +5 -10
- package/src/vector/index.test.ts +687 -192
- package/src/vector/index.ts +57 -29
package/.turbo/turbo-build.log
CHANGED
|
@@ -1,18 +1,23 @@
|
|
|
1
1
|
|
|
2
|
-
> @mastra/chroma@0.1.6-alpha.
|
|
3
|
-
> tsup src/index.ts --format esm --experimental-dts --clean --treeshake
|
|
2
|
+
> @mastra/chroma@0.1.6-alpha.3 build /home/runner/work/mastra/mastra/stores/chroma
|
|
3
|
+
> tsup src/index.ts --format esm,cjs --experimental-dts --clean --treeshake
|
|
4
4
|
|
|
5
5
|
[34mCLI[39m Building entry: src/index.ts
|
|
6
6
|
[34mCLI[39m Using tsconfig: tsconfig.json
|
|
7
7
|
[34mCLI[39m tsup v8.3.6
|
|
8
8
|
[34mTSC[39m Build start
|
|
9
|
-
[32mTSC[39m ⚡️ Build success in
|
|
9
|
+
[32mTSC[39m ⚡️ Build success in 9238ms
|
|
10
10
|
[34mDTS[39m Build start
|
|
11
11
|
[34mCLI[39m Target: es2022
|
|
12
12
|
Analysis will use the bundled TypeScript version 5.7.3
|
|
13
13
|
[36mWriting package typings: /home/runner/work/mastra/mastra/stores/chroma/dist/_tsup-dts-rollup.d.ts[39m
|
|
14
|
-
|
|
14
|
+
Analysis will use the bundled TypeScript version 5.7.3
|
|
15
|
+
[36mWriting package typings: /home/runner/work/mastra/mastra/stores/chroma/dist/_tsup-dts-rollup.d.cts[39m
|
|
16
|
+
[32mDTS[39m ⚡️ Build success in 10295ms
|
|
15
17
|
[34mCLI[39m Cleaning output folder
|
|
16
18
|
[34mESM[39m Build start
|
|
17
|
-
[
|
|
18
|
-
[32mESM[39m
|
|
19
|
+
[34mCJS[39m Build start
|
|
20
|
+
[32mESM[39m [1mdist/index.js [22m[32m7.28 KB[39m
|
|
21
|
+
[32mESM[39m ⚡️ Build success in 715ms
|
|
22
|
+
[32mCJS[39m [1mdist/index.cjs [22m[32m7.31 KB[39m
|
|
23
|
+
[32mCJS[39m ⚡️ Build success in 714ms
|
package/CHANGELOG.md
CHANGED
|
@@ -1,5 +1,42 @@
|
|
|
1
1
|
# @mastra/chroma
|
|
2
2
|
|
|
3
|
+
## 0.1.6-alpha.3
|
|
4
|
+
|
|
5
|
+
### Patch Changes
|
|
6
|
+
|
|
7
|
+
- 0fd78ac: Update vector store functions to use object params
|
|
8
|
+
- fd14a3f: Updating filter location from @mastra/core/filter to @mastra/core/vector/filter
|
|
9
|
+
- 4d4e1e1: Updated vector tests and pinecone
|
|
10
|
+
- bb4f447: Add support for commonjs
|
|
11
|
+
- Updated dependencies [0fd78ac]
|
|
12
|
+
- Updated dependencies [0d25b75]
|
|
13
|
+
- Updated dependencies [fd14a3f]
|
|
14
|
+
- Updated dependencies [3f369a2]
|
|
15
|
+
- Updated dependencies [4d4e1e1]
|
|
16
|
+
- Updated dependencies [bb4f447]
|
|
17
|
+
- @mastra/core@0.4.3-alpha.3
|
|
18
|
+
|
|
19
|
+
## 0.1.6-alpha.2
|
|
20
|
+
|
|
21
|
+
### Patch Changes
|
|
22
|
+
|
|
23
|
+
- Updated dependencies [2512a93]
|
|
24
|
+
- Updated dependencies [e62de74]
|
|
25
|
+
- @mastra/core@0.4.3-alpha.2
|
|
26
|
+
|
|
27
|
+
## 0.1.6-alpha.1
|
|
28
|
+
|
|
29
|
+
### Patch Changes
|
|
30
|
+
|
|
31
|
+
- 5f28f44: Updated Chroma Vector to allow for document storage
|
|
32
|
+
- Updated dependencies [0d185b1]
|
|
33
|
+
- Updated dependencies [ed55f1d]
|
|
34
|
+
- Updated dependencies [8d13b14]
|
|
35
|
+
- Updated dependencies [3ee4831]
|
|
36
|
+
- Updated dependencies [108793c]
|
|
37
|
+
- Updated dependencies [5f28f44]
|
|
38
|
+
- @mastra/core@0.4.3-alpha.1
|
|
39
|
+
|
|
3
40
|
## 0.1.6-alpha.0
|
|
4
41
|
|
|
5
42
|
### Patch Changes
|
package/README.md
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
# @mastra/chroma
|
|
2
2
|
|
|
3
|
-
Vector store implementation for ChromaDB using the official chromadb client with added dimension validation and
|
|
3
|
+
Vector store implementation for ChromaDB using the official chromadb client with added dimension validation, collection management, and document storage capabilities.
|
|
4
4
|
|
|
5
5
|
## Installation
|
|
6
6
|
|
|
@@ -22,21 +22,28 @@ const vectorStore = new ChromaVector({
|
|
|
22
22
|
});
|
|
23
23
|
|
|
24
24
|
// Create a new collection
|
|
25
|
-
await vectorStore.createIndex('my-collection', 1536, 'cosine');
|
|
25
|
+
await vectorStore.createIndex({ indexName: 'my-collection', dimension: 1536, metric: 'cosine' });
|
|
26
26
|
|
|
27
|
-
// Add vectors
|
|
27
|
+
// Add vectors with documents
|
|
28
28
|
const vectors = [[0.1, 0.2, ...], [0.3, 0.4, ...]];
|
|
29
29
|
const metadata = [{ text: 'doc1' }, { text: 'doc2' }];
|
|
30
|
-
const
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
30
|
+
const documents = ['full text 1', 'full text 2'];
|
|
31
|
+
const ids = await vectorStore.upsert({
|
|
32
|
+
indexName: 'my-collection',
|
|
33
|
+
vectors,
|
|
34
|
+
metadata,
|
|
35
|
+
documents, // store original text
|
|
36
|
+
});
|
|
37
|
+
|
|
38
|
+
// Query vectors with document filtering
|
|
39
|
+
const results = await vectorStore.query({
|
|
40
|
+
indexName: 'my-collection',
|
|
41
|
+
queryVector: [0.1, 0.2, ...],
|
|
42
|
+
topK: 10, // topK
|
|
43
|
+
filter: { text: { $eq: 'doc1' } }, // metadata filter
|
|
44
|
+
includeVector: false, // includeVector
|
|
45
|
+
documentFilter: { $contains: 'specific text' } // document content filter
|
|
46
|
+
});
|
|
40
47
|
```
|
|
41
48
|
|
|
42
49
|
## Configuration
|
|
@@ -54,6 +61,8 @@ Optional:
|
|
|
54
61
|
## Features
|
|
55
62
|
|
|
56
63
|
- Vector similarity search with cosine, euclidean, and dot product metrics
|
|
64
|
+
- Document storage and retrieval
|
|
65
|
+
- Document content filtering
|
|
57
66
|
- Strict vector dimension validation
|
|
58
67
|
- Collection-based organization
|
|
59
68
|
- Metadata filtering support
|
|
@@ -64,9 +73,9 @@ Optional:
|
|
|
64
73
|
|
|
65
74
|
## Methods
|
|
66
75
|
|
|
67
|
-
- `createIndex(indexName, dimension, metric?)`: Create a new collection
|
|
68
|
-
- `upsert(indexName, vectors, metadata?, ids?)`: Add or update vectors
|
|
69
|
-
- `query(indexName, queryVector, topK?, filter?, includeVector?)`: Search for similar vectors
|
|
76
|
+
- `createIndex({ indexName, dimension, metric? })`: Create a new collection
|
|
77
|
+
- `upsert({ indexName, vectors, metadata?, ids?, documents? })`: Add or update vectors with optional document storage
|
|
78
|
+
- `query({ indexName, queryVector, topK?, filter?, includeVector?, documentFilter? })`: Search for similar vectors with optional document filtering
|
|
70
79
|
- `listIndexes()`: List all collections
|
|
71
80
|
- `describeIndex(indexName)`: Get collection statistics
|
|
72
81
|
- `deleteIndex(indexName)`: Delete a collection
|
|
@@ -78,6 +87,7 @@ Query results include:
|
|
|
78
87
|
- `id`: Vector ID
|
|
79
88
|
- `score`: Distance/similarity score
|
|
80
89
|
- `metadata`: Associated metadata
|
|
90
|
+
- `document`: Original document text (if stored)
|
|
81
91
|
- `vector`: Original vector (if includeVector is true)
|
|
82
92
|
|
|
83
93
|
## Related Links
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
import { BaseFilterTranslator } from '@mastra/core/vector/filter';
|
|
2
|
+
import type { CreateIndexParams } from '@mastra/core/vector';
|
|
3
|
+
import type { IndexStats } from '@mastra/core/vector';
|
|
4
|
+
import { MastraVector } from '@mastra/core/vector';
|
|
5
|
+
import type { OperatorSupport } from '@mastra/core/vector/filter';
|
|
6
|
+
import type { ParamsToArgs } from '@mastra/core/vector';
|
|
7
|
+
import type { QueryResult } from '@mastra/core/vector';
|
|
8
|
+
import type { QueryVectorParams } from '@mastra/core/vector';
|
|
9
|
+
import type { UpsertVectorParams } from '@mastra/core/vector';
|
|
10
|
+
import type { VectorFilter } from '@mastra/core/vector/filter';
|
|
11
|
+
|
|
12
|
+
/**
|
|
13
|
+
* Translator for Chroma filter queries.
|
|
14
|
+
* Maintains MongoDB-compatible syntax while ensuring proper validation
|
|
15
|
+
* and normalization of values.
|
|
16
|
+
*/
|
|
17
|
+
export declare class ChromaFilterTranslator extends BaseFilterTranslator {
|
|
18
|
+
protected getSupportedOperators(): OperatorSupport;
|
|
19
|
+
translate(filter?: VectorFilter): VectorFilter;
|
|
20
|
+
private translateNode;
|
|
21
|
+
private translateOperator;
|
|
22
|
+
}
|
|
23
|
+
|
|
24
|
+
declare interface ChromaQueryVectorParams extends QueryVectorParams {
|
|
25
|
+
documentFilter?: VectorFilter;
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
declare interface ChromaUpsertVectorParams extends UpsertVectorParams {
|
|
29
|
+
documents?: string[];
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
declare class ChromaVector extends MastraVector {
|
|
33
|
+
private client;
|
|
34
|
+
private collections;
|
|
35
|
+
constructor({ path, auth, }: {
|
|
36
|
+
path: string;
|
|
37
|
+
auth?: {
|
|
38
|
+
provider: string;
|
|
39
|
+
credentials: string;
|
|
40
|
+
};
|
|
41
|
+
});
|
|
42
|
+
getCollection(indexName: string, throwIfNotExists?: boolean): Promise<any>;
|
|
43
|
+
private validateVectorDimensions;
|
|
44
|
+
upsert(...args: ParamsToArgs<ChromaUpsertVectorParams>): Promise<string[]>;
|
|
45
|
+
private HnswSpaceMap;
|
|
46
|
+
createIndex(...args: ParamsToArgs<CreateIndexParams>): Promise<void>;
|
|
47
|
+
transformFilter(filter?: VectorFilter): VectorFilter;
|
|
48
|
+
query(...args: ParamsToArgs<ChromaQueryVectorParams>): Promise<QueryResult[]>;
|
|
49
|
+
listIndexes(): Promise<string[]>;
|
|
50
|
+
describeIndex(indexName: string): Promise<IndexStats>;
|
|
51
|
+
deleteIndex(indexName: string): Promise<void>;
|
|
52
|
+
}
|
|
53
|
+
export { ChromaVector }
|
|
54
|
+
export { ChromaVector as ChromaVector_alias_1 }
|
|
55
|
+
|
|
56
|
+
export { }
|
|
@@ -1,9 +1,13 @@
|
|
|
1
|
-
import { BaseFilterTranslator } from '@mastra/core/filter';
|
|
2
|
-
import {
|
|
3
|
-
import { IndexStats } from '@mastra/core/vector';
|
|
1
|
+
import { BaseFilterTranslator } from '@mastra/core/vector/filter';
|
|
2
|
+
import type { CreateIndexParams } from '@mastra/core/vector';
|
|
3
|
+
import type { IndexStats } from '@mastra/core/vector';
|
|
4
4
|
import { MastraVector } from '@mastra/core/vector';
|
|
5
|
-
import { OperatorSupport } from '@mastra/core/filter';
|
|
6
|
-
import {
|
|
5
|
+
import type { OperatorSupport } from '@mastra/core/vector/filter';
|
|
6
|
+
import type { ParamsToArgs } from '@mastra/core/vector';
|
|
7
|
+
import type { QueryResult } from '@mastra/core/vector';
|
|
8
|
+
import type { QueryVectorParams } from '@mastra/core/vector';
|
|
9
|
+
import type { UpsertVectorParams } from '@mastra/core/vector';
|
|
10
|
+
import type { VectorFilter } from '@mastra/core/vector/filter';
|
|
7
11
|
|
|
8
12
|
/**
|
|
9
13
|
* Translator for Chroma filter queries.
|
|
@@ -12,11 +16,19 @@ import { QueryResult } from '@mastra/core/vector';
|
|
|
12
16
|
*/
|
|
13
17
|
export declare class ChromaFilterTranslator extends BaseFilterTranslator {
|
|
14
18
|
protected getSupportedOperators(): OperatorSupport;
|
|
15
|
-
translate(filter?:
|
|
19
|
+
translate(filter?: VectorFilter): VectorFilter;
|
|
16
20
|
private translateNode;
|
|
17
21
|
private translateOperator;
|
|
18
22
|
}
|
|
19
23
|
|
|
24
|
+
declare interface ChromaQueryVectorParams extends QueryVectorParams {
|
|
25
|
+
documentFilter?: VectorFilter;
|
|
26
|
+
}
|
|
27
|
+
|
|
28
|
+
declare interface ChromaUpsertVectorParams extends UpsertVectorParams {
|
|
29
|
+
documents?: string[];
|
|
30
|
+
}
|
|
31
|
+
|
|
20
32
|
declare class ChromaVector extends MastraVector {
|
|
21
33
|
private client;
|
|
22
34
|
private collections;
|
|
@@ -27,12 +39,13 @@ declare class ChromaVector extends MastraVector {
|
|
|
27
39
|
credentials: string;
|
|
28
40
|
};
|
|
29
41
|
});
|
|
30
|
-
|
|
42
|
+
getCollection(indexName: string, throwIfNotExists?: boolean): Promise<any>;
|
|
31
43
|
private validateVectorDimensions;
|
|
32
|
-
upsert(
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
44
|
+
upsert(...args: ParamsToArgs<ChromaUpsertVectorParams>): Promise<string[]>;
|
|
45
|
+
private HnswSpaceMap;
|
|
46
|
+
createIndex(...args: ParamsToArgs<CreateIndexParams>): Promise<void>;
|
|
47
|
+
transformFilter(filter?: VectorFilter): VectorFilter;
|
|
48
|
+
query(...args: ParamsToArgs<ChromaQueryVectorParams>): Promise<QueryResult[]>;
|
|
36
49
|
listIndexes(): Promise<string[]>;
|
|
37
50
|
describeIndex(indexName: string): Promise<IndexStats>;
|
|
38
51
|
deleteIndex(indexName: string): Promise<void>;
|
package/dist/index.cjs
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
1
|
+
'use strict';
|
|
2
|
+
|
|
3
|
+
var vector = require('@mastra/core/vector');
|
|
4
|
+
var chromadb = require('chromadb');
|
|
5
|
+
var filter = require('@mastra/core/vector/filter');
|
|
6
|
+
|
|
7
|
+
// src/vector/index.ts
|
|
8
|
+
var ChromaFilterTranslator = class extends filter.BaseFilterTranslator {
|
|
9
|
+
getSupportedOperators() {
|
|
10
|
+
return {
|
|
11
|
+
...filter.BaseFilterTranslator.DEFAULT_OPERATORS,
|
|
12
|
+
logical: ["$and", "$or"],
|
|
13
|
+
array: ["$in", "$nin"],
|
|
14
|
+
element: [],
|
|
15
|
+
regex: [],
|
|
16
|
+
custom: []
|
|
17
|
+
};
|
|
18
|
+
}
|
|
19
|
+
translate(filter) {
|
|
20
|
+
if (this.isEmpty(filter)) return filter;
|
|
21
|
+
this.validateFilter(filter);
|
|
22
|
+
return this.translateNode(filter);
|
|
23
|
+
}
|
|
24
|
+
translateNode(node, currentPath = "") {
|
|
25
|
+
if (this.isRegex(node)) {
|
|
26
|
+
throw new Error("Regex is not supported in Chroma");
|
|
27
|
+
}
|
|
28
|
+
if (this.isPrimitive(node)) return this.normalizeComparisonValue(node);
|
|
29
|
+
if (Array.isArray(node)) return { $in: this.normalizeArrayValues(node) };
|
|
30
|
+
const entries = Object.entries(node);
|
|
31
|
+
const firstEntry = entries[0];
|
|
32
|
+
if (entries.length === 1 && firstEntry && this.isOperator(firstEntry[0])) {
|
|
33
|
+
const [operator, value] = firstEntry;
|
|
34
|
+
const translated = this.translateOperator(operator, value);
|
|
35
|
+
return this.isLogicalOperator(operator) ? { [operator]: translated } : translated;
|
|
36
|
+
}
|
|
37
|
+
const result = {};
|
|
38
|
+
const multiOperatorConditions = [];
|
|
39
|
+
for (const [key, value] of entries) {
|
|
40
|
+
const newPath = currentPath ? `${currentPath}.${key}` : key;
|
|
41
|
+
if (this.isOperator(key)) {
|
|
42
|
+
result[key] = this.translateOperator(key, value);
|
|
43
|
+
continue;
|
|
44
|
+
}
|
|
45
|
+
if (typeof value === "object" && value !== null && !Array.isArray(value)) {
|
|
46
|
+
const valueEntries = Object.entries(value);
|
|
47
|
+
if (valueEntries.every(([op]) => this.isOperator(op)) && valueEntries.length > 1) {
|
|
48
|
+
valueEntries.forEach(([op, opValue]) => {
|
|
49
|
+
multiOperatorConditions.push({
|
|
50
|
+
[newPath]: { [op]: this.normalizeComparisonValue(opValue) }
|
|
51
|
+
});
|
|
52
|
+
});
|
|
53
|
+
continue;
|
|
54
|
+
}
|
|
55
|
+
if (Object.keys(value).length === 0) {
|
|
56
|
+
result[newPath] = this.translateNode(value);
|
|
57
|
+
} else {
|
|
58
|
+
const hasOperators = Object.keys(value).some((k) => this.isOperator(k));
|
|
59
|
+
if (hasOperators) {
|
|
60
|
+
const normalizedValue = {};
|
|
61
|
+
for (const [op, opValue] of Object.entries(value)) {
|
|
62
|
+
normalizedValue[op] = this.isOperator(op) ? this.translateOperator(op, opValue) : opValue;
|
|
63
|
+
}
|
|
64
|
+
result[newPath] = normalizedValue;
|
|
65
|
+
} else {
|
|
66
|
+
Object.assign(result, this.translateNode(value, newPath));
|
|
67
|
+
}
|
|
68
|
+
}
|
|
69
|
+
} else {
|
|
70
|
+
result[newPath] = this.translateNode(value);
|
|
71
|
+
}
|
|
72
|
+
}
|
|
73
|
+
if (multiOperatorConditions.length > 0) {
|
|
74
|
+
return { $and: multiOperatorConditions };
|
|
75
|
+
}
|
|
76
|
+
if (Object.keys(result).length > 1 && !currentPath) {
|
|
77
|
+
return {
|
|
78
|
+
$and: Object.entries(result).map(([key, value]) => ({ [key]: value }))
|
|
79
|
+
};
|
|
80
|
+
}
|
|
81
|
+
return result;
|
|
82
|
+
}
|
|
83
|
+
translateOperator(operator, value) {
|
|
84
|
+
if (this.isLogicalOperator(operator)) {
|
|
85
|
+
return Array.isArray(value) ? value.map((item) => this.translateNode(item)) : this.translateNode(value);
|
|
86
|
+
}
|
|
87
|
+
return this.normalizeComparisonValue(value);
|
|
88
|
+
}
|
|
89
|
+
};
|
|
90
|
+
|
|
91
|
+
// src/vector/index.ts
|
|
92
|
+
var ChromaVector = class extends vector.MastraVector {
|
|
93
|
+
client;
|
|
94
|
+
collections;
|
|
95
|
+
constructor({
|
|
96
|
+
path,
|
|
97
|
+
auth
|
|
98
|
+
}) {
|
|
99
|
+
super();
|
|
100
|
+
this.client = new chromadb.ChromaClient({
|
|
101
|
+
path,
|
|
102
|
+
auth
|
|
103
|
+
});
|
|
104
|
+
this.collections = /* @__PURE__ */ new Map();
|
|
105
|
+
}
|
|
106
|
+
async getCollection(indexName, throwIfNotExists = true) {
|
|
107
|
+
try {
|
|
108
|
+
const collection = await this.client.getCollection({ name: indexName, embeddingFunction: void 0 });
|
|
109
|
+
this.collections.set(indexName, collection);
|
|
110
|
+
} catch {
|
|
111
|
+
if (throwIfNotExists) {
|
|
112
|
+
throw new Error(`Index ${indexName} does not exist`);
|
|
113
|
+
}
|
|
114
|
+
return null;
|
|
115
|
+
}
|
|
116
|
+
return this.collections.get(indexName);
|
|
117
|
+
}
|
|
118
|
+
validateVectorDimensions(vectors, dimension) {
|
|
119
|
+
for (let i = 0; i < vectors.length; i++) {
|
|
120
|
+
if (vectors?.[i]?.length !== dimension) {
|
|
121
|
+
throw new Error(
|
|
122
|
+
`Vector at index ${i} has invalid dimension ${vectors?.[i]?.length}. Expected ${dimension} dimensions.`
|
|
123
|
+
);
|
|
124
|
+
}
|
|
125
|
+
}
|
|
126
|
+
}
|
|
127
|
+
async upsert(...args) {
|
|
128
|
+
const params = this.normalizeArgs("upsert", args, ["documents"]);
|
|
129
|
+
const { indexName, vectors, metadata, ids, documents } = params;
|
|
130
|
+
const collection = await this.getCollection(indexName);
|
|
131
|
+
const stats = await this.describeIndex(indexName);
|
|
132
|
+
this.validateVectorDimensions(vectors, stats.dimension);
|
|
133
|
+
const generatedIds = ids || vectors.map(() => crypto.randomUUID());
|
|
134
|
+
const normalizedMetadata = metadata || vectors.map(() => ({}));
|
|
135
|
+
await collection.upsert({
|
|
136
|
+
ids: generatedIds,
|
|
137
|
+
embeddings: vectors,
|
|
138
|
+
metadatas: normalizedMetadata,
|
|
139
|
+
documents
|
|
140
|
+
});
|
|
141
|
+
return generatedIds;
|
|
142
|
+
}
|
|
143
|
+
HnswSpaceMap = {
|
|
144
|
+
cosine: "cosine",
|
|
145
|
+
euclidean: "l2",
|
|
146
|
+
dotproduct: "ip",
|
|
147
|
+
l2: "euclidean",
|
|
148
|
+
ip: "dotproduct"
|
|
149
|
+
};
|
|
150
|
+
async createIndex(...args) {
|
|
151
|
+
const params = this.normalizeArgs("createIndex", args);
|
|
152
|
+
const { indexName, dimension, metric = "cosine" } = params;
|
|
153
|
+
if (!Number.isInteger(dimension) || dimension <= 0) {
|
|
154
|
+
throw new Error("Dimension must be a positive integer");
|
|
155
|
+
}
|
|
156
|
+
const hnswSpace = this.HnswSpaceMap[metric];
|
|
157
|
+
if (!["cosine", "l2", "ip"].includes(hnswSpace)) {
|
|
158
|
+
throw new Error(`Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`);
|
|
159
|
+
}
|
|
160
|
+
await this.client.createCollection({
|
|
161
|
+
name: indexName,
|
|
162
|
+
metadata: {
|
|
163
|
+
dimension,
|
|
164
|
+
"hnsw:space": this.HnswSpaceMap[metric]
|
|
165
|
+
}
|
|
166
|
+
});
|
|
167
|
+
}
|
|
168
|
+
transformFilter(filter) {
|
|
169
|
+
const translator = new ChromaFilterTranslator();
|
|
170
|
+
return translator.translate(filter);
|
|
171
|
+
}
|
|
172
|
+
async query(...args) {
|
|
173
|
+
const params = this.normalizeArgs("query", args, ["documentFilter"]);
|
|
174
|
+
const { indexName, queryVector, topK = 10, filter, includeVector = false, documentFilter } = params;
|
|
175
|
+
const collection = await this.getCollection(indexName, true);
|
|
176
|
+
const defaultInclude = ["documents", "metadatas", "distances"];
|
|
177
|
+
const translatedFilter = this.transformFilter(filter);
|
|
178
|
+
const results = await collection.query({
|
|
179
|
+
queryEmbeddings: [queryVector],
|
|
180
|
+
nResults: topK,
|
|
181
|
+
where: translatedFilter,
|
|
182
|
+
whereDocument: documentFilter,
|
|
183
|
+
include: includeVector ? [...defaultInclude, "embeddings"] : defaultInclude
|
|
184
|
+
});
|
|
185
|
+
return (results.ids[0] || []).map((id, index) => ({
|
|
186
|
+
id,
|
|
187
|
+
score: results.distances?.[0]?.[index] || 0,
|
|
188
|
+
metadata: results.metadatas?.[0]?.[index] || {},
|
|
189
|
+
document: results.documents?.[0]?.[index],
|
|
190
|
+
...includeVector && { vector: results.embeddings?.[0]?.[index] || [] }
|
|
191
|
+
}));
|
|
192
|
+
}
|
|
193
|
+
async listIndexes() {
|
|
194
|
+
const collections = await this.client.listCollections();
|
|
195
|
+
return collections.map((collection) => collection);
|
|
196
|
+
}
|
|
197
|
+
async describeIndex(indexName) {
|
|
198
|
+
const collection = await this.getCollection(indexName);
|
|
199
|
+
const count = await collection.count();
|
|
200
|
+
const metadata = collection.metadata;
|
|
201
|
+
const hnswSpace = metadata?.["hnsw:space"];
|
|
202
|
+
return {
|
|
203
|
+
dimension: metadata?.dimension || 0,
|
|
204
|
+
count,
|
|
205
|
+
metric: this.HnswSpaceMap[hnswSpace]
|
|
206
|
+
};
|
|
207
|
+
}
|
|
208
|
+
async deleteIndex(indexName) {
|
|
209
|
+
await this.client.deleteCollection({ name: indexName });
|
|
210
|
+
this.collections.delete(indexName);
|
|
211
|
+
}
|
|
212
|
+
};
|
|
213
|
+
|
|
214
|
+
exports.ChromaVector = ChromaVector;
|
package/dist/index.d.cts
ADDED
|
@@ -0,0 +1 @@
|
|
|
1
|
+
export { ChromaVector } from './_tsup-dts-rollup.cjs';
|
package/dist/index.js
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
import { BaseFilterTranslator } from '@mastra/core/filter';
|
|
2
1
|
import { MastraVector } from '@mastra/core/vector';
|
|
3
2
|
import { ChromaClient } from 'chromadb';
|
|
3
|
+
import { BaseFilterTranslator } from '@mastra/core/vector/filter';
|
|
4
4
|
|
|
5
5
|
// src/vector/index.ts
|
|
6
6
|
var ChromaFilterTranslator = class extends BaseFilterTranslator {
|
|
@@ -105,7 +105,7 @@ var ChromaVector = class extends MastraVector {
|
|
|
105
105
|
try {
|
|
106
106
|
const collection = await this.client.getCollection({ name: indexName, embeddingFunction: void 0 });
|
|
107
107
|
this.collections.set(indexName, collection);
|
|
108
|
-
} catch
|
|
108
|
+
} catch {
|
|
109
109
|
if (throwIfNotExists) {
|
|
110
110
|
throw new Error(`Index ${indexName} does not exist`);
|
|
111
111
|
}
|
|
@@ -122,7 +122,9 @@ var ChromaVector = class extends MastraVector {
|
|
|
122
122
|
}
|
|
123
123
|
}
|
|
124
124
|
}
|
|
125
|
-
async upsert(
|
|
125
|
+
async upsert(...args) {
|
|
126
|
+
const params = this.normalizeArgs("upsert", args, ["documents"]);
|
|
127
|
+
const { indexName, vectors, metadata, ids, documents } = params;
|
|
126
128
|
const collection = await this.getCollection(indexName);
|
|
127
129
|
const stats = await this.describeIndex(indexName);
|
|
128
130
|
this.validateVectorDimensions(vectors, stats.dimension);
|
|
@@ -131,28 +133,43 @@ var ChromaVector = class extends MastraVector {
|
|
|
131
133
|
await collection.upsert({
|
|
132
134
|
ids: generatedIds,
|
|
133
135
|
embeddings: vectors,
|
|
134
|
-
metadatas: normalizedMetadata
|
|
136
|
+
metadatas: normalizedMetadata,
|
|
137
|
+
documents
|
|
135
138
|
});
|
|
136
139
|
return generatedIds;
|
|
137
140
|
}
|
|
138
|
-
|
|
141
|
+
HnswSpaceMap = {
|
|
142
|
+
cosine: "cosine",
|
|
143
|
+
euclidean: "l2",
|
|
144
|
+
dotproduct: "ip",
|
|
145
|
+
l2: "euclidean",
|
|
146
|
+
ip: "dotproduct"
|
|
147
|
+
};
|
|
148
|
+
async createIndex(...args) {
|
|
149
|
+
const params = this.normalizeArgs("createIndex", args);
|
|
150
|
+
const { indexName, dimension, metric = "cosine" } = params;
|
|
139
151
|
if (!Number.isInteger(dimension) || dimension <= 0) {
|
|
140
152
|
throw new Error("Dimension must be a positive integer");
|
|
141
153
|
}
|
|
154
|
+
const hnswSpace = this.HnswSpaceMap[metric];
|
|
155
|
+
if (!["cosine", "l2", "ip"].includes(hnswSpace)) {
|
|
156
|
+
throw new Error(`Invalid metric: "${metric}". Must be one of: cosine, euclidean, dotproduct`);
|
|
157
|
+
}
|
|
142
158
|
await this.client.createCollection({
|
|
143
159
|
name: indexName,
|
|
144
160
|
metadata: {
|
|
145
161
|
dimension,
|
|
146
|
-
metric
|
|
162
|
+
"hnsw:space": this.HnswSpaceMap[metric]
|
|
147
163
|
}
|
|
148
164
|
});
|
|
149
165
|
}
|
|
150
166
|
transformFilter(filter) {
|
|
151
|
-
const
|
|
152
|
-
|
|
153
|
-
return translatedFilter;
|
|
167
|
+
const translator = new ChromaFilterTranslator();
|
|
168
|
+
return translator.translate(filter);
|
|
154
169
|
}
|
|
155
|
-
async query(
|
|
170
|
+
async query(...args) {
|
|
171
|
+
const params = this.normalizeArgs("query", args, ["documentFilter"]);
|
|
172
|
+
const { indexName, queryVector, topK = 10, filter, includeVector = false, documentFilter } = params;
|
|
156
173
|
const collection = await this.getCollection(indexName, true);
|
|
157
174
|
const defaultInclude = ["documents", "metadatas", "distances"];
|
|
158
175
|
const translatedFilter = this.transformFilter(filter);
|
|
@@ -160,12 +177,14 @@ var ChromaVector = class extends MastraVector {
|
|
|
160
177
|
queryEmbeddings: [queryVector],
|
|
161
178
|
nResults: topK,
|
|
162
179
|
where: translatedFilter,
|
|
180
|
+
whereDocument: documentFilter,
|
|
163
181
|
include: includeVector ? [...defaultInclude, "embeddings"] : defaultInclude
|
|
164
182
|
});
|
|
165
183
|
return (results.ids[0] || []).map((id, index) => ({
|
|
166
184
|
id,
|
|
167
185
|
score: results.distances?.[0]?.[index] || 0,
|
|
168
186
|
metadata: results.metadatas?.[0]?.[index] || {},
|
|
187
|
+
document: results.documents?.[0]?.[index],
|
|
169
188
|
...includeVector && { vector: results.embeddings?.[0]?.[index] || [] }
|
|
170
189
|
}));
|
|
171
190
|
}
|
|
@@ -177,10 +196,11 @@ var ChromaVector = class extends MastraVector {
|
|
|
177
196
|
const collection = await this.getCollection(indexName);
|
|
178
197
|
const count = await collection.count();
|
|
179
198
|
const metadata = collection.metadata;
|
|
199
|
+
const hnswSpace = metadata?.["hnsw:space"];
|
|
180
200
|
return {
|
|
181
201
|
dimension: metadata?.dimension || 0,
|
|
182
202
|
count,
|
|
183
|
-
metric:
|
|
203
|
+
metric: this.HnswSpaceMap[hnswSpace]
|
|
184
204
|
};
|
|
185
205
|
}
|
|
186
206
|
async deleteIndex(indexName) {
|
package/eslint.config.js
ADDED
package/package.json
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
{
|
|
2
2
|
"name": "@mastra/chroma",
|
|
3
|
-
"version": "0.1.6-alpha.
|
|
3
|
+
"version": "0.1.6-alpha.3",
|
|
4
4
|
"description": "Chroma vector store provider for Mastra",
|
|
5
5
|
"type": "module",
|
|
6
6
|
"main": "dist/index.js",
|
|
@@ -10,24 +10,31 @@
|
|
|
10
10
|
"import": {
|
|
11
11
|
"types": "./dist/index.d.ts",
|
|
12
12
|
"default": "./dist/index.js"
|
|
13
|
+
},
|
|
14
|
+
"require": {
|
|
15
|
+
"types": "./dist/index.d.cts",
|
|
16
|
+
"default": "./dist/index.cjs"
|
|
13
17
|
}
|
|
14
18
|
},
|
|
15
19
|
"./package.json": "./package.json"
|
|
16
20
|
},
|
|
17
21
|
"dependencies": {
|
|
18
22
|
"chromadb": "^1.9.4",
|
|
19
|
-
"@mastra/core": "^0.4.3-alpha.
|
|
23
|
+
"@mastra/core": "^0.4.3-alpha.3"
|
|
20
24
|
},
|
|
21
25
|
"devDependencies": {
|
|
22
26
|
"@microsoft/api-extractor": "^7.49.2",
|
|
23
27
|
"@types/node": "^22.13.1",
|
|
24
28
|
"tsup": "^8.0.1",
|
|
25
29
|
"typescript": "^5.7.3",
|
|
26
|
-
"vitest": "^3.0.4"
|
|
30
|
+
"vitest": "^3.0.4",
|
|
31
|
+
"eslint": "^9.20.1",
|
|
32
|
+
"@internal/lint": "0.0.0"
|
|
27
33
|
},
|
|
28
34
|
"scripts": {
|
|
29
|
-
"build": "tsup src/index.ts --format esm --experimental-dts --clean --treeshake",
|
|
35
|
+
"build": "tsup src/index.ts --format esm,cjs --experimental-dts --clean --treeshake",
|
|
30
36
|
"build:watch": "pnpm build --watch",
|
|
31
|
-
"test": "vitest run"
|
|
37
|
+
"test": "vitest run",
|
|
38
|
+
"lint": "eslint ."
|
|
32
39
|
}
|
|
33
40
|
}
|