[ENH] add query config on collection configuration

Author: jairad26

chromadb/api/collection_configuration.py

@@ -99,7 +99,6 @@ def load_collection_configuration_from_json(
                 raise ValueError(
                     f"Could not build embedding function {ef_config['name']} from config {ef_config['config']}: {e}"
                 )
-
     else:
         ef = None
 
@@ -148,11 +147,6 @@ def collection_configuration_to_json(config: CollectionConfiguration) -> Dict[st
     if ef is None:
         ef = None
         ef_config = {"type": "legacy"}
-        return {
-            "hnsw": hnsw_config,
-            "spann": spann_config,
-            "embedding_function": ef_config,
-        }
 
     if ef is not None:
         try:
@@ -260,16 +254,6 @@ class CreateCollectionConfiguration(TypedDict, total=False):
     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"),
-    )
-
-
 def create_collection_configuration_from_legacy_collection_metadata(
     metadata: CollectionMetadata,
 ) -> CreateCollectionConfiguration:
@@ -301,13 +285,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]

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,27 +550,41 @@ 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:
+                return self._embedding_function.embed_query(input=input)
+            else:
+                return self._embedding_function(input=input)
+
         config_ef = self.configuration.get("embedding_function")
         if config_ef is not None:
-            return config_ef(input=input)
+            if is_query:
+                return config_ef.embed_query(input=input)
+            else:
+                return config_ef(input=input)
         if self._embedding_function is None:
             raise ValueError(
                 "You must provide an embedding function to compute embeddings."
                 "https://docs.trychroma.com/guides/embeddings"
             )
-        return self._embedding_function(input=input)
+        if is_query:
+            return self._embedding_function.embed_query(input=input)
+        else:
+            return self._embedding_function(input=input)

chromadb/api/types.py

@@ -569,6 +569,13 @@ class EmbeddingFunction(Protocol[D]):
     def __call__(self, input: D) -> Embeddings:
         ...
 
+    def embed_query(self, input: D) -> Embeddings:
+        """
+        Get the embeddings for a query input.
+        This method is optional, and if not implemented, the default behavior is to call __call__.
+        """
+        return self.__call__(input)
+
     def __init_subclass__(cls) -> None:
         super().__init_subclass__()
         # Raise an exception if __call__ is not defined since it is expected to be defined

chromadb/test/configurations/test_collection_configuration.py

@@ -25,6 +25,7 @@
 from chromadb.test.conftest import ClientFactories
 from chromadb.test.conftest import is_spann_disabled_mode, skip_reason_spann_disabled
 from chromadb.types import Collection as CollectionModel
+from typing import Optional, TypedDict
 
 
 class LegacyEmbeddingFunction(EmbeddingFunction[Embeddable]):
@@ -1616,3 +1617,125 @@ def test_default_space_custom_embedding_function_with_metadata_and_config(
         spann_config = coll.configuration.get("spann")
         assert spann_config is not None
         assert spann_config.get("space") == "ip"
+
+
+class CustomEmbeddingFunctionQueryConfig(TypedDict):
+    task: str
+
+
+@register_embedding_function
+class CustomEmbeddingFunctionWithQueryConfig(EmbeddingFunction[Embeddable]):
+    def __init__(
+        self,
+        task: str,
+        model_name: str,
+        dim: int = 3,
+        query_config: Optional[CustomEmbeddingFunctionQueryConfig] = None,
+    ):
+        self._dim = dim
+        self._model_name = model_name
+        self._task = task
+        self._query_config = query_config
+
+    def __call__(self, input: Embeddable) -> Embeddings:
+        return cast(Embeddings, np.array([[1.0] * self._dim], dtype=np.float32))
+
+    def embed_query(self, input: Embeddable) -> Embeddings:
+        if self._query_config is not None and self._query_config.get("task") == "query":
+            return cast(Embeddings, np.array([[2.0] * self._dim], dtype=np.float32))
+        else:
+            return self.__call__(input)
+
+    @staticmethod
+    def name() -> str:
+        return "custom_ef_with_query_config"
+
+    def get_config(self) -> Dict[str, Any]:
+        return {
+            "model_name": self._model_name,
+            "dim": self._dim,
+            "task": self._task,
+            "query_config": self._query_config,
+        }
+
+    @staticmethod
+    def build_from_config(
+        config: Dict[str, Any]
+    ) -> "CustomEmbeddingFunctionWithQueryConfig":
+        model_name = config.get("model_name")
+        dim = config.get("dim")
+        task = config.get("task")
+        query_config = config.get("query_config")
+
+        if model_name is None or dim is None:
+            assert False, "This code should not be reached"
+
+        return CustomEmbeddingFunctionWithQueryConfig(
+            model_name=model_name, dim=dim, task=task, query_config=query_config  # type: ignore
+        )
+
+    def default_space(self) -> Space:
+        return "cosine"
+
+    def supported_spaces(self) -> List[Space]:
+        return ["cosine"]
+
+
+def test_custom_embedding_function_with_query_config(client: ClientAPI) -> None:
+    client.reset()
+    coll = client.create_collection(
+        name="test_custom_embedding_function_with_query_config",
+        embedding_function=CustomEmbeddingFunctionWithQueryConfig(
+            task="document",
+            model_name="i_want_anything",
+            dim=3,
+            query_config={"task": "query"},
+        ),
+    )
+    assert coll is not None
+    ef = coll.configuration.get("embedding_function")
+    assert ef is not None
+    assert ef.name() == "custom_ef_with_query_config"
+    assert ef.get_config() == {
+        "model_name": "i_want_anything",
+        "dim": 3,
+        "task": "document",
+        "query_config": {"task": "query"},
+    }
+    assert ef.default_space() == "cosine"
+    assert ef.supported_spaces() == ["cosine"]
+    assert np.array_equal(
+        ef.embed_query(input="How many people in Berlin?"),
+        np.array([[2.0, 2.0, 2.0]], dtype=np.float32),
+    )
+
+
+def test_deserializing_custom_embedding_function_with_query_config_no_query_config(
+    client: ClientAPI,
+) -> None:
+    json_string = """
+    {
+        "embedding_function": {
+            "type": "known",
+            "name": "custom_ef_with_query_config",
+            "config": {"model_name": "i_want_anything", "dim": 3, "task": "document"}
+        }
+    }
+    """
+    config = load_collection_configuration_from_json(json.loads(json_string))
+    assert config is not None
+    assert config.get("embedding_function") is not None
+    ef = config.get("embedding_function")
+    assert ef is not None
+    assert ef.get_config() == {
+        "model_name": "i_want_anything",
+        "dim": 3,
+        "task": "document",
+        "query_config": None,
+    }
+    assert ef.default_space() == "cosine"
+    assert ef.supported_spaces() == ["cosine"]
+    assert np.array_equal(
+        ef.embed_query(input="How many people in Berlin?"),
+        np.array([[1.0, 1.0, 1.0]], dtype=np.float32),
+    )

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,
@@ -237,6 +238,7 @@ def config_to_embedding_function(config: Dict[str, Any]) -> EmbeddingFunction:
     "OllamaEmbeddingFunction",
     "InstructorEmbeddingFunction",
     "JinaEmbeddingFunction",
+    "JinaQueryConfig",
     "MistralEmbeddingFunction",
     "MorphEmbeddingFunction",
     "VoyageAIEmbeddingFunction",