Skip to content

Commit 156e21b

Browse files
itaismithitaismith
authored
[ENH] Let Chroma Cloud EFs get API key from client (#5892)
## Description of changes Lets the Chroma Cloud EFs read the API key from the client creating collections configured with them.
1 parent d20abbb commit 156e21b

File tree

19 files changed

+1215
-816
lines changed

19 files changed

+1215
-816
lines changed

‎clients/new-js/packages/ai-embeddings/all/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@chroma-core/all",
3-
"version": "0.1.10",
3+
"version": "0.1.11",
44
"private": false,
55
"description": "All AI embedding providers for Chroma",
66
"main": "dist/cjs/index.cjs",

‎clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@chroma-core/chroma-cloud-qwen",
3-
"version": "0.1.9",
3+
"version": "0.1.10",
44
"private": false,
55
"description": "Chroma embedding function",
66
"main": "dist/cjs/chroma-cloud-qwen.cjs",

‎clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/src/index.test.ts‎

Lines changed: 147 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -1,106 +1,153 @@
11
import {
2-
CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
3-
ChromaCloudQwenEmbeddingFunction,
4-
ChromaCloudQwenEmbeddingModel,
2+
CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
3+
ChromaCloudQwenArgs,
4+
ChromaCloudQwenEmbeddingFunction,
5+
ChromaCloudQwenEmbeddingModel,
56
} from "./index";
67
import { beforeEach, describe, expect, it, jest } from "@jest/globals";
8+
import { CloudClient } from "chromadb";
79

810
describe("ChromaCloudQwenEmbeddingFunction", () => {
9-
beforeEach(() => {
10-
jest.resetAllMocks();
11-
});
12-
13-
const defaultParametersTest = "should initialize with default parameters";
14-
if (!process.env.CHROMA_API_KEY) {
15-
it.skip(defaultParametersTest, () => { });
16-
} else {
17-
it(defaultParametersTest, () => {
18-
const embedder = new ChromaCloudQwenEmbeddingFunction({
19-
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
20-
task: "nl_to_code",
21-
});
22-
expect(embedder.name).toBe("chroma-cloud-qwen");
23-
24-
const config = embedder.getConfig();
25-
expect(config.model).toBe("Qwen/Qwen3-Embedding-0.6B");
26-
expect(config.task).toBe("nl_to_code");
27-
expect(config.api_key_env_var).toBe("CHROMA_API_KEY");
28-
});
29-
}
30-
31-
it("should initialize with custom error for a API key", () => {
32-
const originalEnv = process.env.CHROMA_API_KEY;
33-
delete process.env.CHROMA_API_KEY;
34-
35-
try {
36-
expect(() => {
37-
new ChromaCloudQwenEmbeddingFunction({
38-
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
39-
task: "nl_to_code",
40-
});
41-
}).toThrow("Chroma Embedding API key is required");
42-
} finally {
43-
if (originalEnv) {
44-
process.env.CHROMA_API_KEY = originalEnv;
45-
}
46-
}
47-
});
48-
49-
it("should use custom API key environment variable", () => {
50-
process.env.CUSTOM_CHROMA_API_KEY = "test-api-key";
51-
52-
try {
53-
const embedder = new ChromaCloudQwenEmbeddingFunction({
54-
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
55-
task: "nl_to_code",
56-
apiKeyEnvVar: "CUSTOM_CHROMA_API_KEY",
57-
});
58-
59-
expect(embedder.getConfig().api_key_env_var).toBe(
60-
"CUSTOM_CHROMA_API_KEY",
61-
);
62-
} finally {
63-
delete process.env.CUSTOM_CHROMA_API_KEY;
64-
}
65-
});
66-
67-
const buildFromConfigTest = "should build from config";
68-
if (!process.env.CHROMA_API_KEY) {
69-
it.skip(buildFromConfigTest, () => { });
70-
} else {
71-
it(buildFromConfigTest, () => {
72-
const config = {
73-
api_key_env_var: "CHROMA_API_KEY",
74-
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
75-
instructions: CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
76-
task: "nl_to_code",
77-
};
78-
79-
const embedder = ChromaCloudQwenEmbeddingFunction.buildFromConfig(config);
80-
81-
expect(embedder.getConfig()).toEqual(config);
82-
});
83-
84-
const generateEmbeddingsTest = "should generate embeddings";
85-
if (!process.env.CHROMA_API_KEY) {
86-
it.skip(generateEmbeddingsTest, () => { });
87-
} else {
88-
it(generateEmbeddingsTest, async () => {
89-
const embedder = new ChromaCloudQwenEmbeddingFunction({
90-
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
91-
task: "nl_to_code",
92-
});
93-
const texts = ["Hello world", "Test text"];
94-
const embeddings = await embedder.generate(texts);
95-
96-
expect(embeddings.length).toBe(texts.length);
97-
98-
embeddings.forEach((embedding) => {
99-
expect(embedding.length).toBeGreaterThan(0);
100-
});
101-
102-
expect(embeddings[0]).not.toEqual(embeddings[1]);
103-
});
104-
}
105-
}
11+
beforeEach(() => {
12+
jest.resetAllMocks();
13+
});
14+
15+
const defaultParametersTest = "should initialize with default parameters";
16+
if (!process.env.CHROMA_API_KEY) {
17+
it.skip(defaultParametersTest, () => {});
18+
} else {
19+
it(defaultParametersTest, () => {
20+
const embedder = new ChromaCloudQwenEmbeddingFunction({
21+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
22+
task: "nl_to_code",
23+
});
24+
expect(embedder.name).toBe("chroma-cloud-qwen");
25+
26+
const config = embedder.getConfig();
27+
expect(config.model).toBe("Qwen/Qwen3-Embedding-0.6B");
28+
expect(config.task).toBe("nl_to_code");
29+
expect(config.api_key_env_var).toBe("CHROMA_API_KEY");
30+
});
31+
}
32+
33+
it("should initialize with custom error for a API key", () => {
34+
const originalEnv = process.env.CHROMA_API_KEY;
35+
delete process.env.CHROMA_API_KEY;
36+
37+
try {
38+
expect(() => {
39+
new ChromaCloudQwenEmbeddingFunction({
40+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
41+
task: "nl_to_code",
42+
});
43+
}).toThrow("Chroma Embedding API key is required");
44+
} finally {
45+
if (originalEnv) {
46+
process.env.CHROMA_API_KEY = originalEnv;
47+
}
48+
}
49+
});
50+
51+
it("should use custom API key environment variable", () => {
52+
process.env.CUSTOM_CHROMA_API_KEY = "test-api-key";
53+
54+
try {
55+
const embedder = new ChromaCloudQwenEmbeddingFunction({
56+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
57+
task: "nl_to_code",
58+
apiKeyEnvVar: "CUSTOM_CHROMA_API_KEY",
59+
});
60+
61+
expect(embedder.getConfig().api_key_env_var).toBe(
62+
"CUSTOM_CHROMA_API_KEY",
63+
);
64+
} finally {
65+
delete process.env.CUSTOM_CHROMA_API_KEY;
66+
}
67+
});
68+
69+
const buildFromConfigTest = "should build from config";
70+
if (!process.env.CHROMA_API_KEY) {
71+
it.skip(buildFromConfigTest, () => {});
72+
} else {
73+
it(buildFromConfigTest, () => {
74+
const config = {
75+
api_key_env_var: "CHROMA_API_KEY",
76+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
77+
instructions: CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
78+
task: "nl_to_code",
79+
};
80+
81+
const embedder = ChromaCloudQwenEmbeddingFunction.buildFromConfig(config);
82+
83+
expect(embedder.getConfig()).toEqual(config);
84+
});
85+
86+
const generateEmbeddingsTest = "should generate embeddings";
87+
if (!process.env.CHROMA_API_KEY) {
88+
it.skip(generateEmbeddingsTest, () => {});
89+
} else {
90+
it(generateEmbeddingsTest, async () => {
91+
const embedder = new ChromaCloudQwenEmbeddingFunction({
92+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
93+
task: "nl_to_code",
94+
});
95+
const texts = ["Hello world", "Test text"];
96+
const embeddings = await embedder.generate(texts);
97+
98+
expect(embeddings.length).toBe(texts.length);
99+
100+
embeddings.forEach((embedding) => {
101+
expect(embedding.length).toBeGreaterThan(0);
102+
});
103+
104+
expect(embeddings[0]).not.toEqual(embeddings[1]);
105+
});
106+
}
107+
}
108+
109+
it("should test API key hydration from Chroma Client", async () => {
110+
process.env.CHROMA_API_KEY = "test";
111+
112+
const client = new CloudClient({
113+
tenant: "test-tenant",
114+
database: "test-database",
115+
});
116+
117+
process.env.CHROMA_API_KEY = "";
118+
119+
const config: ChromaCloudQwenArgs = {
120+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
121+
instructions: CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
122+
task: "nl_to_code",
123+
client,
124+
};
125+
126+
const ef = new ChromaCloudQwenEmbeddingFunction(config);
127+
128+
expect(ef).toBeDefined();
129+
expect(ef).toBeInstanceOf(ChromaCloudQwenEmbeddingFunction);
130+
});
131+
132+
it("should test API key hydration from client when building from config", async () => {
133+
const client = new CloudClient({
134+
apiKey: "test",
135+
tenant: "test-tenant",
136+
database: "test-database",
137+
});
138+
139+
const config = {
140+
api_key_env_var: "CHROMA_API_KEY",
141+
model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
142+
instructions: CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
143+
task: "nl_to_code",
144+
};
145+
146+
const embedder = ChromaCloudQwenEmbeddingFunction.buildFromConfig(
147+
config,
148+
client,
149+
);
150+
151+
expect(embedder.getConfig()).toEqual(config);
152+
});
106153
});

0 commit comments

Comments
 (0)