[CLN] Make collection api object wrap model. (#2230)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- This is cleanup borrowed from #1491. It makes Collection.py wrap a
model instead of duplicating data into collection.py and implictly
treating it as a serializable model. This is a cleaner seperation of
concerns. We have the model object, and an API wrapper used to
manipulate it via the api.
 - New functionality
	 - None

## Test plan
*How are these changes tested?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*

Author: HammadB

chromadb/api/fastapi.py

@@ -3,7 +3,6 @@
 from typing import Optional, cast, Tuple
 from typing import Sequence
 from uuid import UUID
-
 import requests
 from overrides import override
 
@@ -41,6 +40,7 @@
 )
 from chromadb.telemetry.product import ProductTelemetryClient
 from urllib.parse import urlparse, urlunparse, quote
+from chromadb.types import Collection as CollectionModel
 
 logger = logging.getLogger(__name__)
 
@@ -209,7 +209,16 @@ def list_collections(
         json_collections = json.loads(resp.text)
         collections = []
         for json_collection in json_collections:
-            collections.append(Collection(self, **json_collection))
+            model = CollectionModel(
+                id=json_collection["id"],
+                name=json_collection["name"],
+                metadata=json_collection["metadata"],
+                dimension=json_collection["dimension"],
+                tenant=json_collection["tenant"],
+                database=json_collection["database"],
+                version=json_collection["version"],
+            )
+            collections.append(Collection(self, model=model))
 
         return collections
 
@@ -254,13 +263,20 @@ def create_collection(
         )
         raise_chroma_error(resp)
         resp_json = json.loads(resp.text)
-        return Collection(
-            client=self,
+        model = CollectionModel(
             id=resp_json["id"],
             name=resp_json["name"],
+            metadata=resp_json["metadata"],
+            dimension=resp_json["dimension"],
+            tenant=resp_json["tenant"],
+            database=resp_json["database"],
+            version=resp_json["version"],
+        )
+        return Collection(
+            client=self,
+            model=model,
             embedding_function=embedding_function,
             data_loader=data_loader,
-            metadata=resp_json["metadata"],
         )
 
     @trace_method("FastAPI.get_collection", OpenTelemetryGranularity.OPERATION)
@@ -288,13 +304,20 @@ def get_collection(
         )
         raise_chroma_error(resp)
         resp_json = json.loads(resp.text)
+        model = CollectionModel(
+            id=resp_json["id"],
+            name=resp_json["name"],
+            metadata=resp_json["metadata"],
+            dimension=resp_json["dimension"],
+            tenant=resp_json["tenant"],
+            database=resp_json["database"],
+            version=resp_json["version"],
+        )
         return Collection(
             client=self,
-            name=resp_json["name"],
-            id=resp_json["id"],
+            model=model,
             embedding_function=embedding_function,
             data_loader=data_loader,
-            metadata=resp_json["metadata"],
         )
 
     @trace_method(

chromadb/api/models/Collection.py

@@ -1,11 +1,7 @@
-from typing import TYPE_CHECKING, Optional, Tuple, Any, Union
-
+from typing import TYPE_CHECKING, Optional, Tuple, Any, Union, cast
 import numpy as np
-from pydantic import BaseModel, PrivateAttr
-
 from uuid import UUID
 import chromadb.utils.embedding_functions as ef
-
 from chromadb.api.types import (
     URI,
     CollectionMetadata,
@@ -46,6 +42,12 @@
     validate_embeddings,
     validate_embedding_function,
 )
+
+# TODO: We should rename the types in chromadb.types to be Models where
+# appropriate. This will help to distinguish between manipulation objects
+# which are essentially API views. And the actual data models which are
+# stored / retrieved / transmitted.
+from chromadb.types import Collection as CollectionModel
 import logging
 
 logger = logging.getLogger(__name__)
@@ -54,33 +56,25 @@
     from chromadb.api import ServerAPI
 
 
-class Collection(BaseModel):
-    name: str
-    id: UUID
-    metadata: Optional[CollectionMetadata] = None
-    tenant: Optional[str] = None
-    database: Optional[str] = None
-    _client: "ServerAPI" = PrivateAttr()
-    _embedding_function: Optional[EmbeddingFunction[Embeddable]] = PrivateAttr()
-    _data_loader: Optional[DataLoader[Loadable]] = PrivateAttr()
+class Collection:
+    _model: CollectionModel
+    _client: "ServerAPI"
+    _embedding_function: Optional[EmbeddingFunction[Embeddable]]
+    _data_loader: Optional[DataLoader[Loadable]]
 
     def __init__(
         self,
         client: "ServerAPI",
-        name: str,
-        id: UUID,
+        model: CollectionModel,
         embedding_function: Optional[
             EmbeddingFunction[Embeddable]
         ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         data_loader: Optional[DataLoader[Loadable]] = None,
-        tenant: Optional[str] = None,
-        database: Optional[str] = None,
-        metadata: Optional[CollectionMetadata] = None,
     ):
-        super().__init__(
-            name=name, metadata=metadata, id=id, tenant=tenant, database=database
-        )
+        """Initializes a new instance of the Collection class."""
+
         self._client = client
+        self._model = model
 
         # Check to make sure the embedding function has the right signature, as defined by the EmbeddingFunction protocol
         if embedding_function is not None:
@@ -92,6 +86,51 @@ def __init__(
     def __repr__(self) -> str:
         return f"Collection(name={self.name})"
 
+    # Expose the model properties as read-only properties on the Collection class
+
+    @property
+    def id(self) -> UUID:
+        return self._model["id"]
+
+    @property
+    def name(self) -> str:
+        return self._model["name"]
+
+    @property
+    def metadata(self) -> CollectionMetadata:
+        return cast(CollectionMetadata, self._model["metadata"])
+
+    @property
+    def tenant(self) -> str:
+        return self._model["tenant"]
+
+    @property
+    def database(self) -> str:
+        return self._model["database"]
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Collection):
+            return False
+        id_match = self.id == other.id
+        name_match = self.name == other.name
+        metadata_match = self.metadata == other.metadata
+        tenant_match = self.tenant == other.tenant
+        database_match = self.database == other.database
+        embedding_function_match = self._embedding_function == other._embedding_function
+        data_loader_match = self._data_loader == other._data_loader
+        return (
+            id_match
+            and name_match
+            and metadata_match
+            and tenant_match
+            and database_match
+            and embedding_function_match
+            and data_loader_match
+        )
+
+    def get_model(self) -> CollectionModel:
+        return self._model
+
     def count(self) -> int:
         """The total number of embeddings added to the database
 
@@ -385,11 +424,14 @@ def modify(
                     "Changing the distance function of a collection once it is created is not supported currently."
                 )
 
+        # Note there is a race condition here where the metadata can be updated
+        # but another thread sees the cached local metadata.
+        # TODO: fixme
         self._client._modify(id=self.id, new_name=name, new_metadata=metadata)
         if name:
-            self.name = name
+            self._model["name"] = name
         if metadata:
-            self.metadata = metadata
+            self._model["metadata"] = metadata
 
     def update(
         self,

chromadb/api/segment.py

@@ -201,13 +201,9 @@ def create_collection(
 
         return Collection(
             client=self,
-            id=coll["id"],
-            name=name,
-            metadata=coll["metadata"],  # type: ignore
+            model=coll,
             embedding_function=embedding_function,
             data_loader=data_loader,
-            tenant=tenant,
-            database=database,
         )
 
     @trace_method(
@@ -260,13 +256,9 @@ def get_collection(
         if existing:
             return Collection(
                 client=self,
-                id=existing[0]["id"],
-                name=existing[0]["name"],
-                metadata=existing[0]["metadata"],  # type: ignore
+                model=existing[0],
                 embedding_function=embedding_function,
                 data_loader=data_loader,
-                tenant=existing[0]["tenant"],
-                database=existing[0]["database"],
             )
         else:
             raise ValueError(f"Collection {name} does not exist.")
@@ -288,11 +280,7 @@ def list_collections(
             collections.append(
                 Collection(
                     client=self,
-                    id=db_collection["id"],
-                    name=db_collection["name"],
-                    metadata=db_collection["metadata"],  # type: ignore
-                    tenant=db_collection["tenant"],
-                    database=db_collection["database"],
+                    model=db_collection,
                 )
             )
         return collections

chromadb/server/fastapi/__init__.py

@@ -1,14 +1,12 @@
 from typing import Any, Callable, cast, Dict, List, Sequence, Optional, Tuple
 import fastapi
 import orjson
-
 from anyio import (
     to_thread,
     CapacityLimiter,
 )
 from fastapi import FastAPI as _FastAPI, Response, Request
 from fastapi.responses import JSONResponse, ORJSONResponse
-
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.routing import APIRoute
 from fastapi import HTTPException, status
@@ -43,9 +41,7 @@
     UpdateEmbedding,
 )
 from starlette.datastructures import Headers
-
 import logging
-
 from chromadb.telemetry.product.events import ServerStartEvent
 from chromadb.utils.fastapi import fastapi_json_response, string_to_uuid as _uuid
 from chromadb.telemetry.opentelemetry.fastapi import instrument_fastapi
@@ -56,6 +52,7 @@
     OpenTelemetryGranularity,
     trace_method,
 )
+from chromadb.types import Collection as CollectionModel
 
 logger = logging.getLogger(__name__)
 
@@ -92,7 +89,7 @@ async def check_http_version_middleware(
     return await call_next(request)
 
 
-class ChromaAPIRouter(fastapi.APIRouter):
+class ChromaAPIRouter(fastapi.APIRouter):  # type: ignore
     # A simple subclass of fastapi's APIRouter which treats URLs with a
     # trailing "/" the same as URLs without. Docs will only contain URLs
     # without trailing "/"s.
@@ -491,7 +488,7 @@ async def list_collections(
         offset: Optional[int] = None,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
-    ) -> Sequence[Collection]:
+    ) -> Sequence[CollectionModel]:
         (
             maybe_tenant,
             maybe_database,
@@ -507,7 +504,7 @@ async def list_collections(
         if maybe_database:
             database = maybe_database
 
-        return cast(
+        api_collections = cast(
             Sequence[Collection],
             await to_thread.run_sync(
                 self._api.list_collections,
@@ -519,6 +516,8 @@ async def list_collections(
             ),
         )
 
+        return [c.get_model() for c in api_collections]
+
     @trace_method("FastAPI.count_collections", OpenTelemetryGranularity.OPERATION)
     async def count_collections(
         self,
@@ -557,7 +556,7 @@ async def create_collection(
         request: Request,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
-    ) -> Collection:
+    ) -> CollectionModel:
         def process_create_collection(
             request: Request, tenant: str, database: str, raw_body: bytes
         ) -> Collection:
@@ -586,7 +585,7 @@ def process_create_collection(
                 database=database,
             )
 
-        return cast(
+        api_collection = cast(
             Collection,
             await to_thread.run_sync(
                 process_create_collection,
@@ -597,6 +596,7 @@ def process_create_collection(
                 limiter=self._capacity_limiter,
             ),
         )
+        return api_collection.get_model()
 
     @trace_method("FastAPI.get_collection", OpenTelemetryGranularity.OPERATION)
     async def get_collection(
@@ -605,7 +605,7 @@ async def get_collection(
         collection_name: str,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
-    ) -> Collection:
+    ) -> CollectionModel:
         (
             maybe_tenant,
             maybe_database,
@@ -621,7 +621,7 @@ async def get_collection(
         if maybe_database:
             database = maybe_database
 
-        return cast(
+        api_collection = cast(
             Collection,
             await to_thread.run_sync(
                 self._api.get_collection,
@@ -634,6 +634,7 @@ async def get_collection(
                 limiter=self._capacity_limiter,
             ),
         )
+        return api_collection.get_model()
 
     @trace_method("FastAPI.update_collection", OpenTelemetryGranularity.OPERATION)
     async def update_collection(