[BUG] Replace empty config objects with None during deserializing

Author: jairad26

chromadb/api/collection_configuration.py

@@ -10,6 +10,9 @@
     known_embedding_functions,
     register_embedding_function,
 )
+from chromadb.utils.embedding_functions.malformed_embedding_function import (
+    MalformedEmbeddingFunction,
+)
 from multiprocessing import cpu_count
 import warnings
 
@@ -83,11 +86,23 @@ def load_collection_configuration_from_json(
         else:
             try:
                 ef = known_embedding_functions[ef_config["name"]]
-                ef = ef.build_from_config(ef_config["config"])  # type: ignore
             except KeyError:
                 raise ValueError(
                     f"Embedding function {ef_config['name']} not found. Add @register_embedding_function decorator to the class definition."
                 )
+            try:
+                ef = ef.build_from_config(ef_config["config"])  # type: ignore
+            except Exception as e:
+                warnings.warn(
+                    f"Could not build embedding function {ef_config['name']} from config {ef_config['config']}: {e} \
+                        Returning a MalformedEmbeddingFunction",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                ef = MalformedEmbeddingFunction(  # type: ignore
+                    malformed_ef_name=ef_config["name"],
+                    config=ef_config["config"],
+                )
     else:
         ef = None
 
@@ -649,7 +664,10 @@ def overwrite_embedding_function(
         return existing_embedding_function
 
     # Validate function compatibility
-    if existing_embedding_function.name() != update_embedding_function.name():
+    if (
+        existing_embedding_function.name() != update_embedding_function.name()
+        and not isinstance(existing_embedding_function, MalformedEmbeddingFunction)
+    ):
         raise ValueError(
             f"Cannot update embedding function: incompatible types "
             f"({existing_embedding_function.name()} vs {update_embedding_function.name()})"

chromadb/test/configurations/test_collection_configuration.py

@@ -16,14 +16,17 @@
     UpdateHNSWConfiguration,
     CreateSpannConfiguration,
     UpdateSpannConfiguration,
-    load_update_collection_configuration_from_json,
     SpannConfiguration,
     overwrite_spann_configuration,
+    overwrite_collection_configuration,
 )
 import json
 import os
 from chromadb.utils.embedding_functions import register_embedding_function
 from chromadb.test.conftest import ClientFactories
+from chromadb.utils.embedding_functions.malformed_embedding_function import (
+    MalformedEmbeddingFunction,
+)
 
 
 # Check if we are running in a mode where SPANN is disabled
@@ -709,19 +712,11 @@ def test_spann_update_from_json(client: ClientAPI) -> None:
         configuration={"spann": initial_spann},
     )
 
-    # Create JSON for update
-    update_json = """
-    {
-        "spann": {
-            "search_nprobe": 15,
-            "ef_search": 200
-        }
-    }
-    """
-
-    # Parse JSON and create update configuration
-    update_config = load_update_collection_configuration_from_json(
-        json.loads(update_json)
+    update_config = UpdateCollectionConfiguration(
+        spann=UpdateSpannConfiguration(
+            search_nprobe=15,
+            ef_search=200,
+        )
     )
 
     # Apply the update
@@ -830,3 +825,68 @@ def test_default_collection_creation(client: ClientAPI) -> None:
     ef = config.get("embedding_function")
     assert ef is not None
     assert ef.name() == "default"
+
+
+def test_malformed_embedding_function() -> None:
+    """Test that on an invalid configuration, the embedding function is a MalformedEmbeddingFunction"""
+    invalid_config: Dict[str, Any] = {
+        "hnsw": {
+            "space": "l2",
+            "ef_construction": 100,
+            "ef_search": 100,
+            "max_neighbors": 16,
+            "resize_factor": 1.2,
+            "sync_threshold": 1000,
+        },
+        "spann": None,
+        "embedding_function": {
+            "name": "custom_ef",
+            "type": "known",
+            "config": {},
+        },
+    }
+    loaded_config = load_collection_configuration_from_json(invalid_config)
+    assert loaded_config is not None
+    loaded_ef = loaded_config.get("embedding_function")
+    assert loaded_ef is not None
+    assert isinstance(loaded_ef, MalformedEmbeddingFunction)
+    assert loaded_ef.malformed_ef_name == "custom_ef"
+
+    assert isinstance(invalid_config["embedding_function"], dict)
+    assert isinstance(invalid_config["embedding_function"]["config"], dict)
+    assert loaded_ef.config == invalid_config["embedding_function"]["config"]
+
+
+def test_update_ef_when_malformed() -> None:
+    invalid_config: Dict[str, Any] = {
+        "hnsw": {
+            "space": "l2",
+            "ef_construction": 100,
+            "ef_search": 100,
+            "max_neighbors": 16,
+            "resize_factor": 1.2,
+            "sync_threshold": 1000,
+        },
+        "spann": None,
+        "embedding_function": {
+            "name": "custom_ef",
+            "type": "known",
+            "config": {},
+        },
+    }
+    loaded_config = load_collection_configuration_from_json(invalid_config)
+    assert loaded_config is not None
+    assert isinstance(
+        loaded_config.get("embedding_function"), MalformedEmbeddingFunction
+    )
+
+    overwrite_config: UpdateCollectionConfiguration = {
+        "embedding_function": CustomEmbeddingFunction(dim=10),
+    }
+
+    new_config = overwrite_collection_configuration(loaded_config, overwrite_config)
+    assert new_config is not None
+    ef = new_config.get("embedding_function")
+    assert ef is not None
+    assert isinstance(ef, CustomEmbeddingFunction)
+    assert ef.get_config() == {"dim": 10}

chromadb/types.py

@@ -8,16 +8,13 @@
 from uuid import UUID
 from enum import Enum
 from pydantic import BaseModel
-import warnings
 
 from chromadb.api.configuration import (
     ConfigurationInternal,
 )
 from chromadb.serde import BaseModelJSONSerializable
 from chromadb.api.collection_configuration import (
     CollectionConfiguration,
-    HNSWConfiguration,
-    SpannConfiguration,
     collection_configuration_to_json,
     load_collection_configuration_from_json,
 )
@@ -149,15 +146,8 @@ def get_configuration(self) -> CollectionConfiguration:
         try:
             return load_collection_configuration_from_json(self.configuration_json)
         except Exception as e:
-            warnings.warn(
-                f"Server does not respond with configuration_json. Please update server: {e}",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-            return CollectionConfiguration(
-                hnsw=HNSWConfiguration(),
-                spann=SpannConfiguration(),
-                embedding_function=None,
+            raise ValueError(
+                f"Could not deserialize configuration_json: {e}",
             )
 
     def set_configuration(self, configuration: CollectionConfiguration) -> None:
@@ -175,19 +165,12 @@ def get_model_fields(self) -> Dict[Any, Any]:
     @override
     def from_json(cls, json_map: Dict[str, Any]) -> Self:
         """Deserializes a Collection object from JSON"""
-        configuration: CollectionConfiguration = {
-            "hnsw": {},
-            "spann": {},
-            "embedding_function": None,
-        }
         try:
             configuration_json = json_map.get("configuration_json", None)
             configuration = load_collection_configuration_from_json(configuration_json)
         except Exception as e:
-            warnings.warn(
-                f"Server does not respond with configuration_json. Please update server: {e}",
-                DeprecationWarning,
-                stacklevel=2,
+            raise ValueError(
+                f"Could not deserialize configuration_json: {e}",
             )
         return cls(
             id=json_map["id"],

chromadb/utils/embedding_functions/baseten_embedding_function.py

@@ -1,16 +1,18 @@
 import os
-from chromadb.utils.embedding_functions.openai_embedding_function import OpenAIEmbeddingFunction
-from chromadb.api.types import Documents, EmbeddingFunction
+from chromadb.utils.embedding_functions.openai_embedding_function import (
+    OpenAIEmbeddingFunction,
+)
 from typing import Dict, Any, Optional
+import warnings
 
-class BasetenEmbeddingFunction(OpenAIEmbeddingFunction):
 
+class BasetenEmbeddingFunction(OpenAIEmbeddingFunction):
     def __init__(
-            self,
-            api_key: Optional[str],
-            api_base: str,
-            api_key_env_var: str = "CHROMA_BASETEN_API_KEY",
-            ):
+        self,
+        api_key: Optional[str],
+        api_base: str,
+        api_key_env_var: str = "CHROMA_BASETEN_API_KEY",
+    ):
         """
         Initialize the BasetenEmbeddingFunction.
         Args:
@@ -25,33 +27,35 @@ def __init__(
                 "The openai python package is not installed. Please install it with `pip install openai`"
             )
 
+        if api_key is not None:
+            warnings.warn(
+                "Direct api_key configuration will not be persisted. "
+                "Please use environment variables via api_key_env_var for persistent storage.",
+                DeprecationWarning,
+            )
+
         self.api_key_env_var = api_key_env_var
         # Prioritize api_key argument, then environment variable
         resolved_api_key = api_key or os.getenv(api_key_env_var)
         if not resolved_api_key:
-            raise ValueError(f"API key not provided and {api_key_env_var} environment variable is not set.")
+            raise ValueError(
+                f"API key not provided and {api_key_env_var} environment variable is not set."
+            )
         self.api_key = resolved_api_key
         if not api_base:
             raise ValueError("The api_base argument must be provided.")
         self.api_base = api_base
         self.model_name = "baseten-embedding-model"
         self.dimensions = None
 
-        self.client = openai.OpenAI(
-            api_key=self.api_key,
-            base_url=self.api_base
-        )
-        
+        self.client = openai.OpenAI(api_key=self.api_key, base_url=self.api_base)
+
     @staticmethod
     def name() -> str:
         return "baseten"
-    
+
     def get_config(self) -> Dict[str, Any]:
-        return {
-            "api_base": self.api_base,
-            "api_key_env_var": self.api_key_env_var
-        }
-        
+        return {"api_base": self.api_base, "api_key_env_var": self.api_key_env_var}
 
     @staticmethod
     def build_from_config(config: Dict[str, Any]) -> "BasetenEmbeddingFunction":
@@ -68,16 +72,20 @@ def build_from_config(config: Dict[str, Any]) -> "BasetenEmbeddingFunction":
         api_key_env_var = config.get("api_key_env_var")
         api_base = config.get("api_base")
         if api_key_env_var is None or api_base is None:
-            raise ValueError("Missing 'api_key_env_var' or 'api_base' in configuration for BasetenEmbeddingFunction.")
+            raise ValueError(
+                "Missing 'api_key_env_var' or 'api_base' in configuration for BasetenEmbeddingFunction."
+            )
 
         # Note: We rely on the __init__ method to handle potential missing api_key
         # by checking the environment variable if the config value is None.
         # However, api_base must be present either in config or have a default.
         if api_base is None:
-             raise ValueError("Missing 'api_base' in configuration for BasetenEmbeddingFunction.")
+            raise ValueError(
+                "Missing 'api_base' in configuration for BasetenEmbeddingFunction."
+            )
 
         return BasetenEmbeddingFunction(
-            api_key=None, # Pass None if not in config, __init__ will check env var
+            api_key=None,  # Pass None if not in config, __init__ will check env var
             api_base=api_base,
             api_key_env_var=api_key_env_var,
-        )
+        )

chromadb/utils/embedding_functions/cloudflare_workers_ai_embedding_function.py

@@ -8,6 +8,7 @@
 import os
 from chromadb.utils.embedding_functions.schemas import validate_config_schema
 from typing import cast
+import warnings
 
 BASE_URL = "https://api.cloudflare.com/client/v4/accounts"
 GATEWAY_BASE_URL = "https://gateway.ai.cloudflare.com/v1"
@@ -43,6 +44,13 @@ def __init__(
             raise ValueError(
                 "The httpx python package is not installed. Please install it with `pip install httpx`"
             )
+
+        if api_key is not None:
+            warnings.warn(
+                "Direct api_key configuration will not be persisted. "
+                "Please use environment variables via api_key_env_var for persistent storage.",
+                DeprecationWarning,
+            )
         self.model_name = model_name
         self.account_id = account_id
         self.api_key_env_var = api_key_env_var

chromadb/utils/embedding_functions/cohere_embedding_function.py

@@ -13,6 +13,7 @@
 import base64
 import io
 import importlib
+import warnings
 
 
 class CohereEmbeddingFunction(EmbeddingFunction[Embeddable]):
@@ -36,6 +37,12 @@ def __init__(
                 "The PIL python package is not installed. Please install it with `pip install pillow`"
             )
 
+        if api_key is not None:
+            warnings.warn(
+                "Direct api_key configuration will not be persisted. "
+                "Please use environment variables via api_key_env_var for persistent storage.",
+                DeprecationWarning,
+            )
         self.api_key_env_var = api_key_env_var
         self.api_key = api_key or os.getenv(api_key_env_var)
         if not self.api_key: