[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

@@ -576,6 +576,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/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",