[ENH] Add schema support to collection configuration

Author: jairad26

chromadb/api/collection_configuration.py

@@ -8,6 +8,7 @@
     EmbeddingFunction,
     QueryConfig,
 )
+from chromadb.base_types import CollectionSchema
 from chromadb.utils.embedding_functions import (
     known_embedding_functions,
     register_embedding_function,
@@ -44,6 +45,7 @@ class CollectionConfiguration(TypedDict, total=True):
     spann: Optional[SpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_embedding_function: Optional[EmbeddingFunction]  # type: ignore
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def load_collection_configuration_from_json_str(
@@ -126,6 +128,7 @@ def load_collection_configuration_from_json(
         spann=spann_config,
         embedding_function=ef,  # type: ignore
         query_embedding_function=query_ef,  # type: ignore
+        schema=config_json_map.get("schema"),
     )
 
 
@@ -278,6 +281,7 @@ class CreateCollectionConfiguration(TypedDict, total=False):
     spann: Optional[CreateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_config: Optional[QueryConfig]
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def create_collection_configuration_from_legacy_collection_metadata(
@@ -416,6 +420,7 @@ def create_collection_configuration_to_json(
         "spann": spann_config,
         "embedding_function": ef_config,
         "query_config": query_config,
+        "schema": config.get("schema"),
     }
 
 
@@ -488,6 +493,7 @@ class UpdateCollectionConfiguration(TypedDict, total=False):
     spann: Optional[UpdateSpannConfiguration]
     embedding_function: Optional[EmbeddingFunction]  # type: ignore
     query_config: Optional[QueryConfig]
+    schema: Optional[Dict[str, CollectionSchema]]
 
 
 def update_collection_configuration_from_legacy_collection_metadata(
@@ -587,6 +593,7 @@ def update_collection_configuration_to_json(
         "spann": spann_config,
         "embedding_function": ef_config,
         "query_config": query_config,
+        "schema": config.get("schema"),
     }
 
 
@@ -750,14 +757,34 @@ def overwrite_collection_configuration(
                     ef_config[k] = v
                 query_ef = updated_embedding_function.build_from_config(ef_config)
 
+    existing_schema = existing_config.get("schema")
+    new_diff_schema = update_config.get("schema")
+    updated_schema: Optional[Dict[str, CollectionSchema]] = None
+    if existing_schema is not None:
+        if new_diff_schema is not None:
+            updated_schema = overwrite_schema(existing_schema, new_diff_schema)
+        else:
+            updated_schema = existing_schema
+    else:
+        updated_schema = new_diff_schema
+
     return CollectionConfiguration(
         hnsw=updated_hnsw_config,
         spann=updated_spann_config,
         embedding_function=updated_embedding_function,
         query_embedding_function=query_ef,
+        schema=updated_schema,
     )
 
 
+def overwrite_schema(
+    existing_schema: Dict[str, CollectionSchema],
+    new_diff_schema: Dict[str, CollectionSchema],
+) -> Dict[str, CollectionSchema]:
+    """Overwrite a schema with a new configuration"""
+    return {**existing_schema, **new_diff_schema}
+
+
 def validate_embedding_function_conflict_on_create(
     embedding_function: Optional[EmbeddingFunction],  # type: ignore
     configuration_ef: Optional[EmbeddingFunction],  # type: ignore

chromadb/api/models/AsyncCollection.py

@@ -60,6 +60,10 @@ async def add(
             ValueError: If you provide an id that already exists
 
         """
+
+        # get metadatas, if they don't exist in schema yet add them, and do collection.modify()
+        # also validate that the metadatas are valid for the schema
+
         add_request = self._validate_and_prepare_add_request(
             ids=ids,
             embeddings=embeddings,

chromadb/api/models/Collection.py

@@ -77,15 +77,23 @@ def add(
 
         """
 
-        add_request = self._validate_and_prepare_add_request(
+        curr_schema = self._model.get_configuration().get("schema")
+        add_request, new_attributes = self._validate_and_prepare_add_request(
             ids=ids,
             embeddings=embeddings,
             metadatas=metadatas,
             documents=documents,
             images=images,
             uris=uris,
+            schema=curr_schema,
         )
 
+        if len(new_attributes.keys()) > 0:
+            if curr_schema is None:
+                curr_schema = {}
+            curr_schema = {**curr_schema, **new_attributes}
+            self.modify(configuration={"schema": curr_schema})
+
         self._client._add(
             collection_id=self.id,
             ids=add_request["ids"],
@@ -317,7 +325,8 @@ def update(
         Returns:
             None
         """
-        update_request = self._validate_and_prepare_update_request(
+        curr_schema = self._model.get_configuration().get("schema")
+        update_request, new_attributes = self._validate_and_prepare_update_request(
             ids=ids,
             embeddings=embeddings,
             metadatas=metadatas,
@@ -326,6 +335,12 @@ def update(
             uris=uris,
         )
 
+        if len(new_attributes.keys()) > 0:
+            if curr_schema is None:
+                curr_schema = {}
+            curr_schema = {**curr_schema, **new_attributes}
+            self.modify(configuration={"schema": curr_schema})
+
         self._client._update(
             collection_id=self.id,
             ids=update_request["ids"],
@@ -362,7 +377,8 @@ def upsert(
         Returns:
             None
         """
-        upsert_request = self._validate_and_prepare_upsert_request(
+        curr_schema = self._model.get_configuration().get("schema")
+        upsert_request, new_attributes = self._validate_and_prepare_upsert_request(
             ids=ids,
             embeddings=embeddings,
             metadatas=metadatas,
@@ -371,6 +387,12 @@ def upsert(
             uris=uris,
         )
 
+        if len(new_attributes.keys()) > 0:
+            if curr_schema is None:
+                curr_schema = {}
+            curr_schema = {**curr_schema, **new_attributes}
+            self.modify(configuration={"schema": curr_schema})
+
         self._client._upsert(
             collection_id=self.id,
             ids=upsert_request["ids"],

chromadb/api/models/CollectionCommon.py

@@ -10,12 +10,14 @@
     TypeVar,
     Union,
     cast,
+    Tuple,
 )
 from chromadb.types import Metadata
 import numpy as np
 from uuid import UUID
-
 import chromadb.utils.embedding_functions as ef
+from chromadb.base_types import CollectionSchema
+
 from chromadb.api.types import (
     URI,
     URIs,
@@ -47,6 +49,7 @@
     maybe_cast_one_to_many,
     normalize_base_record_set,
     normalize_insert_record_set,
+    update_schema_with_insert_record_set,
     validate_base_record_set,
     validate_ids,
     validate_include,
@@ -204,7 +207,8 @@ def _validate_and_prepare_add_request(
         documents: Optional[OneOrMany[Document]],
         images: Optional[OneOrMany[Image]],
         uris: Optional[OneOrMany[URI]],
-    ) -> AddRequest:
+        schema: Optional[Dict[str, CollectionSchema]],
+    ) -> Tuple[AddRequest, Dict[str, CollectionSchema]]:
         # Unpack
         add_records = normalize_insert_record_set(
             ids=ids,
@@ -216,22 +220,29 @@ def _validate_and_prepare_add_request(
         )
 
         # Validate
-        validate_insert_record_set(record_set=add_records)
+        validate_insert_record_set(record_set=add_records, schema=schema)
         validate_record_set_contains_any(record_set=add_records, contains_any={"ids"})
 
+        new_attributes = update_schema_with_insert_record_set(
+            record_set=add_records, schema=schema
+        )
+
         # Prepare
         if add_records["embeddings"] is None:
             validate_record_set_for_embedding(record_set=add_records)
             add_embeddings = self._embed_record_set(record_set=add_records)
         else:
             add_embeddings = add_records["embeddings"]
 
-        return AddRequest(
-            ids=add_records["ids"],
-            embeddings=add_embeddings,
-            metadatas=add_records["metadatas"],
-            documents=add_records["documents"],
-            uris=add_records["uris"],
+        return (
+            AddRequest(
+                ids=add_records["ids"],
+                embeddings=add_embeddings,
+                metadatas=add_records["metadatas"],
+                documents=add_records["documents"],
+                uris=add_records["uris"],
+            ),
+            new_attributes,
         )
 
     @validation_context("get")
@@ -350,7 +361,8 @@ def _validate_and_prepare_update_request(
         documents: Optional[OneOrMany[Document]],
         images: Optional[OneOrMany[Image]],
         uris: Optional[OneOrMany[URI]],
-    ) -> UpdateRequest:
+        schema: Optional[Dict[str, CollectionSchema]],
+    ) -> Tuple[UpdateRequest, Dict[str, CollectionSchema]]:
         # Unpack
         update_records = normalize_insert_record_set(
             ids=ids,
@@ -362,7 +374,10 @@ def _validate_and_prepare_update_request(
         )
 
         # Validate
-        validate_insert_record_set(record_set=update_records)
+        validate_insert_record_set(record_set=update_records, schema=schema)
+        new_attributes = update_schema_with_insert_record_set(
+            record_set=update_records, schema=schema
+        )
 
         # Prepare
         if update_records["embeddings"] is None:
@@ -380,12 +395,15 @@ def _validate_and_prepare_update_request(
         else:
             update_embeddings = update_records["embeddings"]
 
-        return UpdateRequest(
-            ids=update_records["ids"],
-            embeddings=update_embeddings,
-            metadatas=update_records["metadatas"],
-            documents=update_records["documents"],
-            uris=update_records["uris"],
+        return (
+            UpdateRequest(
+                ids=update_records["ids"],
+                embeddings=update_embeddings,
+                metadatas=update_records["metadatas"],
+                documents=update_records["documents"],
+                uris=update_records["uris"],
+            ),
+            new_attributes,
         )
 
     @validation_context("upsert")
@@ -402,7 +420,8 @@ def _validate_and_prepare_upsert_request(
         documents: Optional[OneOrMany[Document]] = None,
         images: Optional[OneOrMany[Image]] = None,
         uris: Optional[OneOrMany[URI]] = None,
-    ) -> UpsertRequest:
+        schema: Optional[Dict[str, CollectionSchema]] = None,
+    ) -> Tuple[UpsertRequest, Dict[str, CollectionSchema]]:
         # Unpack
         upsert_records = normalize_insert_record_set(
             ids=ids,
@@ -414,8 +433,10 @@ def _validate_and_prepare_upsert_request(
         )
 
         # Validate
-        validate_insert_record_set(record_set=upsert_records)
-
+        validate_insert_record_set(record_set=upsert_records, schema=schema)
+        new_attributes = update_schema_with_insert_record_set(
+            record_set=upsert_records, schema=schema
+        )
         # Prepare
         if upsert_records["embeddings"] is None:
             validate_record_set_for_embedding(
@@ -425,12 +446,15 @@ def _validate_and_prepare_upsert_request(
         else:
             upsert_embeddings = upsert_records["embeddings"]
 
-        return UpsertRequest(
-            ids=upsert_records["ids"],
-            metadatas=upsert_records["metadatas"],
-            embeddings=upsert_embeddings,
-            documents=upsert_records["documents"],
-            uris=upsert_records["uris"],
+        return (
+            UpsertRequest(
+                ids=upsert_records["ids"],
+                metadatas=upsert_records["metadatas"],
+                embeddings=upsert_embeddings,
+                documents=upsert_records["documents"],
+                uris=upsert_records["uris"],
+            ),
+            new_attributes,
         )
 
     @validation_context("delete")