[ENH] add query config on collection configuration

jairad26 · jairad26 · commit 69696fde285d · 2025-06-20T17:17:04.000-07:00
diff --git a/chromadb/api/collection_configuration.py b/chromadb/api/collection_configuration.py
@@ -1,10 +1,12 @@
 from typing import TypedDict, Dict, Any, Optional, cast, get_args
 import json
+import copy
 from chromadb.api.types import (
     Space,
     CollectionMetadata,
     UpdateMetadata,
     EmbeddingFunction,
+    QueryConfig,
 )
 from chromadb.utils.embedding_functions import (
     known_embedding_functions,
@@ -41,6 +43,7 @@ class CollectionConfiguration(TypedDict, total=True):
     hnsw: Optional[HNSWConfiguration]
     spann: Optional[SpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
+    query_embedding_function: Optional[EmbeddingFunction]  # type: ignore
 
 
 def load_collection_configuration_from_json_str(
@@ -64,6 +67,8 @@ def load_collection_configuration_from_json(
     spann_config = None
     ef_config = None
 
+    query_ef = None
+
     # Process vector index configuration (HNSW or SPANN)
     if config_json_map.get("hnsw") is not None:
         hnsw_config = cast(HNSWConfiguration, config_json_map["hnsw"])
@@ -100,13 +105,27 @@ def load_collection_configuration_from_json(
                     f"Could not build embedding function {ef_config['name']} from config {ef_config['config']}: {e}"
                 )
 
+            if config_json_map.get("query_config") is not None:
+                query_config = config_json_map["query_config"]
+                query_ef_config = copy.deepcopy(ef_config)
+                query_ef = known_embedding_functions[ef_name]
+                for k, v in query_config.items():
+                    query_ef_config["config"][k] = v
+
+                try:
+                    query_ef = query_ef.build_from_config(query_ef_config["config"])  # type: ignore
+                except Exception as e:
+                    raise ValueError(
+                        f"Could not build query embedding function {query_ef_config['name']} from config {query_ef_config['config']}: {e}"
+                    )
     else:
         ef = None
 
     return CollectionConfiguration(
         hnsw=hnsw_config,
         spann=spann_config,
         embedding_function=ef,  # type: ignore
+        query_embedding_function=query_ef,  # type: ignore
     )
 
 
@@ -258,16 +277,7 @@ class CreateCollectionConfiguration(TypedDict, total=False):
     hnsw: Optional[CreateHNSWConfiguration]
     spann: Optional[CreateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
-
-
-def load_collection_configuration_from_create_collection_configuration(
-    config: CreateCollectionConfiguration,
-) -> CollectionConfiguration:
-    return CollectionConfiguration(
-        hnsw=config.get("hnsw"),
-        spann=config.get("spann"),
-        embedding_function=config.get("embedding_function"),
-    )
+    query_config: Optional[QueryConfig]
 
 
 def create_collection_configuration_from_legacy_collection_metadata(
@@ -301,13 +311,6 @@ def create_collection_configuration_from_legacy_metadata_dict(
     return CreateCollectionConfiguration(hnsw=hnsw_config)
 
 
-def load_create_collection_configuration_from_json_str(
-    json_str: str,
-) -> CreateCollectionConfiguration:
-    json_map = json.loads(json_str)
-    return load_create_collection_configuration_from_json(json_map)
-
-
 # TODO: make warnings prettier and add link to migration docs
 def load_create_collection_configuration_from_json(
     json_map: Dict[str, Any]
@@ -353,6 +356,7 @@ def create_collection_configuration_to_json(
 ) -> Dict[str, Any]:
     """Convert a CreateCollection configuration to a JSON-serializable dict"""
     ef_config: Dict[str, Any] | None = None
+    query_config: Dict[str, Any] | None = None
     hnsw_config = config.get("hnsw")
     spann_config = config.get("spann")
     if hnsw_config is not None:
@@ -389,6 +393,15 @@ def create_collection_configuration_to_json(
                 "config": ef.get_config(),
             }
             register_embedding_function(type(ef))  # type: ignore
+
+            q = config.get("query_config")
+            if q is not None:
+                if q.name() == ef.name():
+                    query_config = q.get_config()
+                else:
+                    raise ValueError(
+                        f"query config name {q.name()} does not match embedding function name {ef.name()}"
+                    )
     except Exception as e:
         warnings.warn(
             f"legacy embedding function config: {e}",
@@ -402,6 +415,7 @@ def create_collection_configuration_to_json(
         "hnsw": hnsw_config,
         "spann": spann_config,
         "embedding_function": ef_config,
+        "query_config": query_config,
     }
 
 
@@ -473,6 +487,7 @@ class UpdateCollectionConfiguration(TypedDict, total=False):
     hnsw: Optional[UpdateHNSWConfiguration]
     spann: Optional[UpdateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
+    query_config: Optional[QueryConfig]
 
 
 def update_collection_configuration_from_legacy_collection_metadata(
@@ -528,6 +543,7 @@ def update_collection_configuration_to_json(
     hnsw_config = config.get("hnsw")
     spann_config = config.get("spann")
     ef = config.get("embedding_function")
+    query_config: Dict[str, Any] | None = None
     if hnsw_config is None and spann_config is None and ef is None:
         return {}
 
@@ -555,13 +571,22 @@ def update_collection_configuration_to_json(
                 "config": ef.get_config(),
             }
             register_embedding_function(type(ef))  # type: ignore
+            q = config.get("query_config")
+            if q is not None:
+                if q.name() == ef.name():
+                    query_config = q.get_config()
+                else:
+                    raise ValueError(
+                        f"query config name {q.name()} does not match embedding function name {ef.name()}"
+                    )
     else:
         ef_config = None
 
     return {
         "hnsw": hnsw_config,
         "spann": spann_config,
         "embedding_function": ef_config,
+        "query_config": query_config,
     }
 
 
@@ -710,10 +735,26 @@ def overwrite_collection_configuration(
         else:
             updated_embedding_function = update_ef
 
+    query_ef = None
+    if updated_embedding_function is not None:
+        q = update_config.get("query_config")
+        if q is not None:
+            if q.name() != updated_embedding_function.name():
+                raise ValueError(
+                    f"query config name {q.name()} does not match embedding function name {updated_embedding_function.name()}"
+                )
+            else:
+                ef_config = copy.deepcopy(updated_embedding_function.get_config())
+                query_config = q.get_config()
+                for k, v in query_config.items():
+                    ef_config[k] = v
+                query_ef = updated_embedding_function.build_from_config(ef_config)
+
     return CollectionConfiguration(
         hnsw=updated_hnsw_config,
         spann=updated_spann_config,
         embedding_function=updated_embedding_function,
+        query_embedding_function=query_ef,
     )
 
 
diff --git a/chromadb/api/models/CollectionCommon.py b/chromadb/api/models/CollectionCommon.py
@@ -313,7 +313,9 @@ def _validate_and_prepare_query_request(
         # Prepare
         if query_records["embeddings"] is None:
             validate_record_set_for_embedding(record_set=query_records)
-            request_embeddings = self._embed_record_set(record_set=query_records)
+            request_embeddings = self._embed_record_set(
+                record_set=query_records, is_query=True
+            )
         else:
             request_embeddings = query_records["embeddings"]
 
@@ -531,7 +533,10 @@ def _update_model_after_modify_success(
             )
 
     def _embed_record_set(
-        self, record_set: BaseRecordSet, embeddable_fields: Optional[Set[str]] = None
+        self,
+        record_set: BaseRecordSet,
+        embeddable_fields: Optional[Set[str]] = None,
+        is_query: bool = False,
     ) -> Embeddings:
         if embeddable_fields is None:
             embeddable_fields = get_default_embeddable_record_set_fields()
@@ -545,21 +550,30 @@ def _embed_record_set(
                             "You must set a data loader on the collection if loading from URIs."
                         )
                     return self._embed(
-                        input=self._data_loader(uris=cast(URIs, record_set[field]))  # type: ignore[literal-required]
+                        input=self._data_loader(uris=cast(URIs, record_set[field])),  # type: ignore[literal-required]
+                        is_query=is_query,
                     )
                 else:
-                    return self._embed(input=record_set[field])  # type: ignore[literal-required]
+                    return self._embed(
+                        input=record_set[field],  # type: ignore[literal-required]
+                        is_query=is_query,
+                    )
         raise ValueError(
             "Record does not contain any non-None fields that can be embedded."
             f"Embeddable Fields: {embeddable_fields}"
             f"Record Fields: {record_set}"
         )
 
-    def _embed(self, input: Any) -> Embeddings:
+    def _embed(self, input: Any, is_query: bool = False) -> Embeddings:
         if self._embedding_function is not None and not isinstance(
             self._embedding_function, ef.DefaultEmbeddingFunction
         ):
             return self._embedding_function(input=input)
+        if is_query:
+            config_ef = self.configuration.get("query_embedding_function")
+            if config_ef is not None:
+                return config_ef(input=input)
+
         config_ef = self.configuration.get("embedding_function")
         if config_ef is not None:
             return config_ef(input=input)
diff --git a/chromadb/api/types.py b/chromadb/api/types.py
@@ -672,6 +672,16 @@ def is_legacy(self) -> bool:
         return False
 
 
+class QueryConfig:
+    @abstractmethod
+    def name(self) -> str:
+        return NotImplemented
+
+    @abstractmethod
+    def get_config(self) -> Dict[str, Any]:
+        return NotImplemented
+
+
 def validate_embedding_function(
     embedding_function: EmbeddingFunction[Embeddable],
 ) -> None:
diff --git a/chromadb/utils/embedding_functions/__init__.py b/chromadb/utils/embedding_functions/__init__.py
@@ -32,6 +32,7 @@
 )
 from chromadb.utils.embedding_functions.jina_embedding_function import (
     JinaEmbeddingFunction,
+    JinaQueryConfig,
 )
 from chromadb.utils.embedding_functions.voyageai_embedding_function import (
     VoyageAIEmbeddingFunction,
@@ -232,6 +233,7 @@ def config_to_embedding_function(config: Dict[str, Any]) -> EmbeddingFunction:
     "OllamaEmbeddingFunction",
     "InstructorEmbeddingFunction",
     "JinaEmbeddingFunction",
+    "JinaQueryConfig",
     "MistralEmbeddingFunction",
     "VoyageAIEmbeddingFunction",
     "ONNXMiniLM_L6_V2",
diff --git a/chromadb/utils/embedding_functions/jina_embedding_function.py b/chromadb/utils/embedding_functions/jina_embedding_function.py
@@ -1,6 +1,13 @@
-from chromadb.api.types import Embeddings, Documents, EmbeddingFunction, Space
+from chromadb.api.types import (
+    Embeddings,
+    Documents,
+    EmbeddingFunction,
+    Space,
+    QueryConfig,
+)
 from chromadb.utils.embedding_functions.schemas import validate_config_schema
 from typing import List, Dict, Any, Union, Optional
+from typing_extensions import override
 import os
 import numpy as np
 import warnings
@@ -206,3 +213,17 @@ def validate_config(config: Dict[str, Any]) -> None:
             ValidationError: If the configuration does not match the schema
         """
         validate_config_schema(config, "jina")
+
+
+class JinaQueryConfig(QueryConfig):
+    def __init__(self, task: Optional[str] = None):
+        self.task = task
+
+    @override
+    def name(self) -> str:
+        return "jina"
+
+    def get_config(self) -> Dict[str, Any]:
+        return {
+            "task": self.task,
+        }
diff --git a/go/pkg/sysdb/coordinator/model/collection_configuration.go b/go/pkg/sysdb/coordinator/model/collection_configuration.go
@@ -56,6 +56,7 @@ type SpannConfiguration struct {
 type InternalCollectionConfiguration struct {
 	VectorIndex       *VectorIndexConfiguration       `json:"vector_index"`
 	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
+	QueryConfig       interface{}                     `json:"query_config,omitempty"`
 }
 
 // DefaultHnswCollectionConfiguration returns a default configuration using HNSW
@@ -127,4 +128,5 @@ type UpdateVectorIndexConfiguration struct {
 type InternalUpdateCollectionConfiguration struct {
 	VectorIndex       *UpdateVectorIndexConfiguration `json:"vector_index,omitempty"`
 	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
+	QueryConfig       interface{}                     `json:"query_config,omitempty"`
 }
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -854,6 +854,10 @@ func (tc *Catalog) updateCollectionConfiguration(
 		existingConfig.EmbeddingFunction = updateConfig.EmbeddingFunction
 	}
 
+	if updateConfig.QueryConfig != nil {
+		existingConfig.QueryConfig = updateConfig.QueryConfig
+	}
+
 	// Serialize updated config back to JSON
 	updatedConfigBytes, err := json.Marshal(existingConfig)
 	if err != nil {
diff --git a/rust/python_bindings/src/bindings.rs b/rust/python_bindings/src/bindings.rs
@@ -280,6 +280,7 @@ impl Bindings {
                     hnsw: None,
                     spann: None,
                     embedding_function: None,
+                    query_config: None,
                 },
                 self.frontend.get_default_knn_index(),
             )?),
diff --git a/rust/segment/src/distributed_hnsw.rs b/rust/segment/src/distributed_hnsw.rs
@@ -420,6 +420,7 @@ pub mod test {
                     },
                 ),
                 embedding_function: None,
+                query_config: None,
             },
             ..Default::default()
         };
diff --git a/rust/segment/src/distributed_spann.rs b/rust/segment/src/distributed_spann.rs
@@ -629,6 +629,7 @@ mod test {
             config: chroma_types::InternalCollectionConfiguration {
                 vector_index: chroma_types::VectorIndexConfiguration::Spann(params),
                 embedding_function: None,
+                query_config: None,
             },
             metadata: None,
             dimension: None,
@@ -845,6 +846,7 @@ mod test {
             config: InternalCollectionConfiguration {
                 vector_index: chroma_types::VectorIndexConfiguration::Spann(params),
                 embedding_function: None,
+                query_config: None,
             },
             ..Default::default()
         };
diff --git a/rust/types/src/collection_configuration.rs b/rust/types/src/collection_configuration.rs

Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,7 @@ type SpannConfiguration struct {`
`56`	`56`	`type InternalCollectionConfiguration struct {`
`57`	`57`	VectorIndex *VectorIndexConfiguration `json:"vector_index"`
`58`	`58`	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
	`59`	+ QueryConfig interface{} `json:"query_config,omitempty"`
`59`	`60`	`}`
`60`	`61`
`61`	`62`	`// DefaultHnswCollectionConfiguration returns a default configuration using HNSW`
`@@ -127,4 +128,5 @@ type UpdateVectorIndexConfiguration struct {`
`127`	`128`	`type InternalUpdateCollectionConfiguration struct {`
`128`	`129`	VectorIndex *UpdateVectorIndexConfiguration `json:"vector_index,omitempty"`
`129`	`130`	EmbeddingFunction *EmbeddingFunctionConfiguration `json:"embedding_function,omitempty"`
	`131`	+ QueryConfig interface{} `json:"query_config,omitempty"`
`130`	`132`	`}`