[ENH] Add support for default space in create coll config

Author: jairad26

chromadb/api/async_fastapi.py

@@ -297,7 +297,7 @@ async def create_collection(
     ) -> CollectionModel:
         """Creates a collection"""
         config_json = (
-            create_collection_configuration_to_json(configuration)
+            create_collection_configuration_to_json(configuration, metadata)
             if configuration
             else None
         )

chromadb/api/collection_configuration.py

@@ -12,6 +12,7 @@
 )
 from multiprocessing import cpu_count
 import warnings
+from chromadb.errors import InvalidSpaceError
 
 
 class HNSWConfiguration(TypedDict, total=False):
@@ -342,14 +343,16 @@ def load_create_collection_configuration_from_json(
 
 def create_collection_configuration_to_json_str(
     config: CreateCollectionConfiguration,
+    metadata: Optional[CollectionMetadata] = None,
 ) -> str:
     """Convert a CreateCollection configuration to a JSON-serializable string"""
-    return json.dumps(create_collection_configuration_to_json(config))
+    return json.dumps(create_collection_configuration_to_json(config, metadata))
 
 
 # TODO: make warnings prettier and add link to migration docs
 def create_collection_configuration_to_json(
     config: CreateCollectionConfiguration,
+    metadata: Optional[CollectionMetadata] = None,
 ) -> Dict[str, Any]:
     """Convert a CreateCollection configuration to a JSON-serializable dict"""
     ef_config: Dict[str, Any] | None = None
@@ -383,12 +386,47 @@ def create_collection_configuration_to_json(
         if ef.is_legacy():
             ef_config = {"type": "legacy"}
         else:
+            if hnsw_config is None and spann_config is None:
+                if metadata is None or metadata.get("hnsw:space") is None:
+                    # this populates space from ef if not provided in either config
+                    hnsw_config = CreateHNSWConfiguration(space=ef.default_space())
+
+            # if hnsw config or spann config exists but space is not provided, populate it from ef
+            if hnsw_config is not None and hnsw_config.get("space") is None:
+                hnsw_config["space"] = ef.default_space()
+            if spann_config is not None and spann_config.get("space") is None:
+                spann_config["space"] = ef.default_space()
+
+            # Validate space compatibility with embedding function
+            if hnsw_config is not None:
+                if hnsw_config.get("space") not in ef.supported_spaces():
+                    raise InvalidSpaceError(
+                        f"space {hnsw_config.get('space')} is not supported by {ef.name()}. Supported spaces: {ef.supported_spaces()}"
+                    )
+            if spann_config is not None:
+                if spann_config.get("space") not in ef.supported_spaces():
+                    raise InvalidSpaceError(
+                        f"space {spann_config.get('space')} is not supported by {ef.name()}. Supported spaces: {ef.supported_spaces()}"
+                    )
+            if (
+                hnsw_config is None
+                and spann_config is None
+                and metadata is not None
+                and metadata.get("hnsw:space") is not None
+            ):
+                if metadata.get("hnsw:space") not in ef.supported_spaces():
+                    raise InvalidSpaceError(
+                        f"space {metadata.get('hnsw:space')} is not supported by {ef.name()}. Supported spaces: {ef.supported_spaces()}"
+                    )
+
             ef_config = {
                 "name": ef.name(),
                 "type": "known",
                 "config": ef.get_config(),
             }
             register_embedding_function(type(ef))  # type: ignore
+    except InvalidSpaceError:
+        raise
     except Exception as e:
         warnings.warn(
             f"legacy embedding function config: {e}",

chromadb/api/fastapi.py

@@ -258,7 +258,9 @@ def create_collection(
             json={
                 "name": name,
                 "metadata": metadata,
-                "configuration": create_collection_configuration_to_json(configuration)
+                "configuration": create_collection_configuration_to_json(
+                    configuration, metadata
+                )
                 if configuration
                 else None,
                 "get_or_create": get_or_create,

chromadb/api/rust.py

@@ -219,7 +219,7 @@ def create_collection(
         )
         if configuration:
             configuration_json_str = create_collection_configuration_to_json_str(
-                configuration
+                configuration, metadata
             )
         else:
             configuration_json_str = None

chromadb/api/segment.py

@@ -236,7 +236,7 @@ def create_collection(
             name=name,
             metadata=metadata,
             configuration_json=create_collection_configuration_to_json(
-                configuration or CreateCollectionConfiguration()
+                configuration or CreateCollectionConfiguration(), metadata
             ),
             tenant=tenant,
             database=database,

chromadb/db/impl/grpc/client.py

@@ -6,6 +6,7 @@
     create_collection_configuration_to_json_str,
     UpdateCollectionConfiguration,
     update_collection_configuration_to_json_str,
+    CollectionMetadata,
 )
 from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System, logger
 from chromadb.db.system import SysDB
@@ -327,7 +328,7 @@ def create_collection(
                 id=id.hex,
                 name=name,
                 configuration_json_str=create_collection_configuration_to_json_str(
-                    configuration
+                    configuration, cast(CollectionMetadata, metadata)
                 ),
                 metadata=to_proto_update_metadata(metadata) if metadata else None,
                 dimension=dimension,

chromadb/db/mixins/sysdb.py

@@ -43,6 +43,7 @@
     collection_configuration_to_json_str,
     overwrite_collection_configuration,
     update_collection_configuration_from_legacy_update_metadata,
+    CollectionMetadata,
 )
 
 logger = logging.getLogger(__name__)
@@ -311,7 +312,9 @@ def create_collection(
         collection = Collection(
             id=id,
             name=name,
-            configuration_json=create_collection_configuration_to_json(configuration),
+            configuration_json=create_collection_configuration_to_json(
+                configuration, cast(CollectionMetadata, metadata)
+            ),
             metadata=metadata,
             dimension=dimension,
             tenant=tenant,
@@ -337,7 +340,9 @@ def create_collection(
                     ParameterValue(self.uuid_to_db(collection["id"])),
                     ParameterValue(collection["name"]),
                     ParameterValue(
-                        create_collection_configuration_to_json_str(configuration)
+                        create_collection_configuration_to_json_str(
+                            configuration, cast(CollectionMetadata, metadata)
+                        )
                     ),
                     ParameterValue(collection["dimension"]),
                     # Get the database id for the database with the given name and tenant
@@ -941,7 +946,7 @@ def _insert_config_from_legacy_params(
         create_collection_config = CreateCollectionConfiguration()
         # Write the configuration into the database
         configuration_json_str = create_collection_configuration_to_json_str(
-            create_collection_config
+            create_collection_config, cast(CollectionMetadata, metadata)
         )
         q = (
             self.querybuilder()

chromadb/errors.py

@@ -173,6 +173,17 @@ def name(cls) -> str:
         return "QuotaError"
 
 
+class InvalidSpaceError(ChromaError):
+    @overrides
+    def code(self) -> int:
+        return 400
+
+    @classmethod
+    @overrides
+    def name(cls) -> str:
+        return "InvalidSpaceError"
+
+
 error_types: Dict[str, Type[ChromaError]] = {
     "InvalidDimension": InvalidDimensionException,
     "InvalidArgumentError": InvalidArgumentError,
@@ -189,6 +200,7 @@ def name(cls) -> str:
     "UniqueConstraintError": UniqueConstraintError,
     "QuotaError": QuotaError,
     "InternalError": InternalError,
+    "InvalidSpaceError": InvalidSpaceError,
     # Catch-all for any other errors
     "ChromaError": ChromaError,
 }