[ENH]: fix Qwen EF hydration with custom prompts/tasks (#5808)

codetheweb · web-flow · commit 4bc9153ab6bf · 2025-11-04T23:28:07.000Z
diff --git a/chromadb/test/utils/test_embedding_function_schemas.py b/chromadb/test/utils/test_embedding_function_schemas.py
@@ -50,13 +50,13 @@ def test_embedding_function_config_roundtrip(
         if ef_name == "chroma-cloud-qwen":
             from chromadb.utils.embedding_functions.chroma_cloud_qwen_embedding_function import (
                 ChromaCloudQwenEmbeddingModel,
-                ChromaCloudQwenEmbeddingTask,
                 CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
             )
+
             mock_ef.get_config.return_value = {
                 "api_key_env_var": "CHROMA_API_KEY",
                 "model": ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B.value,
-                "task": ChromaCloudQwenEmbeddingTask.NL_TO_CODE.value,
+                "task": "nl_to_code",
                 "instructions": CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
             }
 
diff --git a/chromadb/utils/embedding_functions/chroma_cloud_qwen_embedding_function.py b/chromadb/utils/embedding_functions/chroma_cloud_qwen_embedding_function.py
@@ -5,31 +5,34 @@
 from chromadb.utils.embedding_functions.schemas import validate_config_schema
 from enum import Enum
 
+
 class ChromaCloudQwenEmbeddingModel(Enum):
     QWEN3_EMBEDDING_0p6B = "Qwen/Qwen3-Embedding-0.6B"
 
-class ChromaCloudQwenEmbeddingTask(Enum):
-    NL_TO_CODE = "nl_to_code"
 
 class ChromaCloudQwenEmbeddingTarget(Enum):
     DOCUMENTS = "documents"
     QUERY = "query"
 
-ChromaCloudQwenEmbeddingInstructions = Dict[ChromaCloudQwenEmbeddingTask, Dict[ChromaCloudQwenEmbeddingTarget, str]]
+
+ChromaCloudQwenEmbeddingInstructions = Dict[
+    str, Dict[ChromaCloudQwenEmbeddingTarget, str]
+]
 
 CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS: ChromaCloudQwenEmbeddingInstructions = {
-    ChromaCloudQwenEmbeddingTask.NL_TO_CODE: {
+    "nl_to_code": {
         ChromaCloudQwenEmbeddingTarget.DOCUMENTS: "",
         # Taken from https://github.com/QwenLM/Qwen3-Embedding/blob/main/evaluation/task_prompts.json
         ChromaCloudQwenEmbeddingTarget.QUERY: "Given a question about coding, retrieval code or passage that can solve user's question",
-    }        
+    }
 }
 
+
 class ChromaCloudQwenEmbeddingFunction(EmbeddingFunction[Documents]):
     def __init__(
         self,
         model: ChromaCloudQwenEmbeddingModel,
-        task: ChromaCloudQwenEmbeddingTask,
+        task: str,
         instructions: ChromaCloudQwenEmbeddingInstructions = CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
         api_key_env_var: str = "CHROMA_API_KEY",
     ):
@@ -38,11 +41,11 @@ def __init__(
 
         Args:
             model (ChromaCloudQwenEmbeddingModel): The specific Qwen model to use for embeddings.
-            task (ChromaCloudQwenEmbeddingTask): The task for which embeddings are being generated.
+            task (str): The task for which embeddings are being generated.
             instructions (ChromaCloudQwenEmbeddingInstructions, optional): A dictionary containing
                 custom instructions to use for the specified Qwen model. Defaults to CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS.
             api_key_env_var (str, optional): Environment variable name that contains your API key.
-                Defaults to "CHROMA_API_KEY".                
+                Defaults to "CHROMA_API_KEY".
         """
         try:
             import httpx
@@ -63,7 +66,10 @@ def __init__(
         self._api_url = "https://embed.trychroma.com"
         self._session = httpx.Client()
         self._session.headers.update(
-            {"x-chroma-token": self.api_key, "x-chroma-embedding-model": self.model.value}
+            {
+                "x-chroma-token": self.api_key,
+                "x-chroma-embedding-model": self.model.value,
+            }
         )
 
     def _parse_response(self, response: Any) -> Embeddings:
@@ -83,7 +89,6 @@ def _parse_response(self, response: Any) -> Embeddings:
 
         return np.array(embeddings, dtype=np.float32)
 
-
     def __call__(self, input: Documents) -> Embeddings:
         """
         Generate embeddings for the given documents.
@@ -98,7 +103,9 @@ def __call__(self, input: Documents) -> Embeddings:
             return []
 
         payload: Dict[str, Union[str, Documents]] = {
-            "instructions": self.instructions[self.task][ChromaCloudQwenEmbeddingTarget.DOCUMENTS],
+            "instructions": self.instructions[self.task][
+                ChromaCloudQwenEmbeddingTarget.DOCUMENTS
+            ],
             "texts": input,
         }
 
@@ -114,7 +121,9 @@ def embed_query(self, input: Documents) -> Embeddings:
             return []
 
         payload: Dict[str, Union[str, Documents]] = {
-            "instructions": self.instructions[self.task][ChromaCloudQwenEmbeddingTarget.QUERY],
+            "instructions": self.instructions[self.task][
+                ChromaCloudQwenEmbeddingTarget.QUERY
+            ],
             "texts": input,
         }
 
@@ -147,34 +156,30 @@ def build_from_config(config: Dict[str, Any]) -> "EmbeddingFunction[Documents]":
         if instructions is not None:
             deserialized_instructions = {}
             for task_key, targets in instructions.items():
-                # Convert string key to enum
-                task_enum = ChromaCloudQwenEmbeddingTask(task_key)
-                deserialized_instructions[task_enum] = {}
+                deserialized_instructions[task_key] = {}
                 for target_key, instruction in targets.items():
                     # Convert string key to enum
                     target_enum = ChromaCloudQwenEmbeddingTarget(target_key)
-                    deserialized_instructions[task_enum][target_enum] = instruction
+                    deserialized_instructions[task_key][target_enum] = instruction
+                    deserialized_instructions[task_key][target_enum] = instruction
 
         return ChromaCloudQwenEmbeddingFunction(
             model=ChromaCloudQwenEmbeddingModel(model),
-            task=ChromaCloudQwenEmbeddingTask(task),
+            task=task,
             instructions=deserialized_instructions,
             api_key_env_var=api_key_env_var or "CHROMA_API_KEY",
         )
 
     def get_config(self) -> Dict[str, Any]:
         # Serialize instructions dict with enum keys to string keys for JSON compatibility
         serialized_instructions = {
-            task.value: {
-                target.value: instruction
-                for target, instruction in targets.items()
-            }
+            task: {target.value: instruction for target, instruction in targets.items()}
             for task, targets in self.instructions.items()
         }
         return {
             "api_key_env_var": self.api_key_env_var,
             "model": self.model.value,
-            "task": self.task.value,
+            "task": self.task,
             "instructions": serialized_instructions,
         }
 
@@ -192,7 +197,7 @@ def validate_config_update(
         elif "instructions" in new_config:
             raise ValueError(
                 "The instructions cannot be changed after the embedding function has been initialized."
-            )        
+            )
 
     @staticmethod
     def validate_config(config: Dict[str, Any]) -> None:
@@ -205,4 +210,4 @@ def validate_config(config: Dict[str, Any]) -> None:
         Raises:
             ValidationError: If the configuration does not match the schema
         """
-        validate_config_schema(config, "chroma-cloud-qwen")
+        validate_config_schema(config, "chroma-cloud-qwen")
diff --git a/clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/src/index.test.ts b/clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/src/index.test.ts
@@ -2,7 +2,6 @@ import {
 	CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
 	ChromaCloudQwenEmbeddingFunction,
 	ChromaCloudQwenEmbeddingModel,
-	ChromaCloudQwenEmbeddingTask,
 } from "./index";
 import { beforeEach, describe, expect, it, jest } from "@jest/globals";
 
@@ -13,12 +12,11 @@ describe("ChromaCloudQwenEmbeddingFunction", () => {
 
 	const defaultParametersTest = "should initialize with default parameters";
 	if (!process.env.CHROMA_API_KEY) {
-		it.skip(defaultParametersTest, () => {});
+		it.skip(defaultParametersTest, () => { });
 	} else {
 		it(defaultParametersTest, () => {
 			const embedder = new ChromaCloudQwenEmbeddingFunction({
 				model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
-				task: ChromaCloudQwenEmbeddingTask.NL_TO_CODE,
 			});
 			expect(embedder.name).toBe("chroma-cloud-qwen");
 
@@ -37,7 +35,6 @@ describe("ChromaCloudQwenEmbeddingFunction", () => {
 			expect(() => {
 				new ChromaCloudQwenEmbeddingFunction({
 					model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
-					task: ChromaCloudQwenEmbeddingTask.NL_TO_CODE,
 				});
 			}).toThrow("Chroma Embedding API key is required");
 		} finally {
@@ -53,7 +50,6 @@ describe("ChromaCloudQwenEmbeddingFunction", () => {
 		try {
 			const embedder = new ChromaCloudQwenEmbeddingFunction({
 				model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
-				task: ChromaCloudQwenEmbeddingTask.NL_TO_CODE,
 				apiKeyEnvVar: "CUSTOM_CHROMA_API_KEY",
 			});
 
@@ -67,14 +63,14 @@ describe("ChromaCloudQwenEmbeddingFunction", () => {
 
 	const buildFromConfigTest = "should build from config";
 	if (!process.env.CHROMA_API_KEY) {
-		it.skip(buildFromConfigTest, () => {});
+		it.skip(buildFromConfigTest, () => { });
 	} else {
 		it(buildFromConfigTest, () => {
 			const config = {
 				api_key_env_var: "CHROMA_API_KEY",
 				model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
-				task: ChromaCloudQwenEmbeddingTask.NL_TO_CODE,
 				instructions: CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
+				task: "nl_to_code",
 			};
 
 			const embedder = ChromaCloudQwenEmbeddingFunction.buildFromConfig(config);
@@ -84,12 +80,11 @@ describe("ChromaCloudQwenEmbeddingFunction", () => {
 
 		const generateEmbeddingsTest = "should generate embeddings";
 		if (!process.env.CHROMA_API_KEY) {
-			it.skip(generateEmbeddingsTest, () => {});
+			it.skip(generateEmbeddingsTest, () => { });
 		} else {
 			it(generateEmbeddingsTest, async () => {
 				const embedder = new ChromaCloudQwenEmbeddingFunction({
 					model: ChromaCloudQwenEmbeddingModel.QWEN3_EMBEDDING_0p6B,
-					task: ChromaCloudQwenEmbeddingTask.NL_TO_CODE,
 				});
 				const texts = ["Hello world", "Test text"];
 				const embeddings = await embedder.generate(texts);
diff --git a/clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/src/index.ts b/clients/new-js/packages/ai-embeddings/chroma-cloud-qwen/src/index.ts
@@ -11,9 +11,10 @@ import {
 
 const NAME = "chroma-cloud-qwen";
 
+
 export interface ChromaCloudQwenConfig {
 	model: ChromaCloudQwenEmbeddingModel;
-	task: ChromaCloudQwenEmbeddingTask;
+	task: string;
 	instructions: ChromaCloudQwenEmbeddingInstructions;
 	api_key_env_var: string;
 }
@@ -22,33 +23,29 @@ export enum ChromaCloudQwenEmbeddingModel {
 	QWEN3_EMBEDDING_0p6B = "Qwen/Qwen3-Embedding-0.6B",
 }
 
-export enum ChromaCloudQwenEmbeddingTask {
-	NL_TO_CODE = "nl_to_code", // Queries are in natural language, documents are code
-}
-
 export enum ChromaCloudQwenEmbeddingTarget {
 	DOCUMENTS = "documents",
 	QUERY = "query",
 }
 
 export type ChromaCloudQwenEmbeddingInstructions = Record<
-	ChromaCloudQwenEmbeddingTask,
+	string,
 	Record<ChromaCloudQwenEmbeddingTarget, string>
 >;
 
 export const CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS: ChromaCloudQwenEmbeddingInstructions =
-	{
-		[ChromaCloudQwenEmbeddingTask.NL_TO_CODE]: {
-			[ChromaCloudQwenEmbeddingTarget.DOCUMENTS]: "",
-			[ChromaCloudQwenEmbeddingTarget.QUERY]:
-				// Taken from https://github.com/QwenLM/Qwen3-Embedding/blob/main/evaluation/task_prompts.json
-				"Given a question about coding, retrieval code or passage that can solve user's question",
-		},
-	};
+{
+	"nl_to_code": {
+		[ChromaCloudQwenEmbeddingTarget.DOCUMENTS]: "",
+		[ChromaCloudQwenEmbeddingTarget.QUERY]:
+			// Taken from https://github.com/QwenLM/Qwen3-Embedding/blob/main/evaluation/task_prompts.json
+			"Given a question about coding, retrieval code or passage that can solve user's question",
+	},
+};
 
 export interface ChromaCloudQwenArgs {
 	model: ChromaCloudQwenEmbeddingModel;
-	task: ChromaCloudQwenEmbeddingTask;
+	task?: string;
 	instructions?: ChromaCloudQwenEmbeddingInstructions;
 	apiKeyEnvVar?: string;
 }
@@ -70,13 +67,13 @@ export class ChromaCloudQwenEmbeddingFunction implements EmbeddingFunction {
 	private readonly model: ChromaCloudQwenEmbeddingModel;
 	private readonly url: string;
 	private readonly headers: { [key: string]: string };
-	private readonly task: ChromaCloudQwenEmbeddingTask;
+	private readonly task: string;
 	private readonly instructions: ChromaCloudQwenEmbeddingInstructions;
 
 	constructor(args: ChromaCloudQwenArgs) {
 		const {
 			model,
-			task,
+			task = "nl_to_code",
 			instructions = CHROMA_CLOUD_QWEN_DEFAULT_INSTRUCTIONS,
 			apiKeyEnvVar = "CHROMA_API_KEY",
 		} = args;
@@ -185,16 +182,14 @@ export class ChromaCloudQwenEmbeddingFunction implements EmbeddingFunction {
 		if (config.instructions) {
 			deserializedInstructions = {} as ChromaCloudQwenEmbeddingInstructions;
 			for (const [taskKey, targets] of Object.entries(config.instructions)) {
-				// taskKey is the enum value string like "nl_to_code"
-				const taskEnum = taskKey as ChromaCloudQwenEmbeddingTask;
-				deserializedInstructions[taskEnum] = {} as Record<
+				deserializedInstructions[taskKey] = {} as Record<
 					ChromaCloudQwenEmbeddingTarget,
 					string
 				>;
 				for (const [targetKey, instruction] of Object.entries(targets)) {
 					// targetKey is the enum value string like "documents" or "query"
 					const targetEnum = targetKey as ChromaCloudQwenEmbeddingTarget;
-					deserializedInstructions[taskEnum][targetEnum] = instruction;
+					deserializedInstructions[taskKey][targetEnum] = instruction;
 				}
 			}
 		} else {
