Improved component model

.pre-commit-config.yaml

@@ -34,5 +34,5 @@ repos:
     rev: 'v1.2.0'
     hooks:
     -   id: mypy
-        args: [--strict, --ignore-missing-imports, --follow-imports=silent]
+        args: [--strict, --ignore-missing-imports, --follow-imports=silent, --disable-error-code=type-abstract]
         additional_dependencies: ["types-requests", "pydantic", "overrides", "hypothesis", "pytest"]

.vscode/settings.json

@@ -33,6 +33,7 @@
     "--ignore-missing-imports",
     "--show-column-numbers",
     "--no-pretty",
-    "--strict"
+    "--strict",
+    "--disable-error-code=type-abstract"
   ]
 }

chromadb/__init__.py

@@ -1,6 +1,7 @@
 import chromadb.config
 import logging
 from chromadb.telemetry.events import ClientStartEvent
+from chromadb.telemetry import Telemetry
 from chromadb.config import Settings, System
 from chromadb.api import API
 
@@ -22,14 +23,16 @@ def get_settings() -> Settings:
 
 
 def Client(settings: Settings = __settings) -> API:
-    """Return a chroma.API instance based on the provided or environmental
-    settings, optionally overriding the DB instance."""
+    """Return a running chroma.API instance"""
 
     system = System(settings)
 
-    telemetry_client = system.get_telemetry()
+    telemetry_client = system.instance(Telemetry)
+    api = system.instance(API)
+
+    system.start()
 
     # Submit event for client start
     telemetry_client.capture(ClientStartEvent())
 
-    return system.get_api()
+    return api

chromadb/api/__init__.py

@@ -16,15 +16,10 @@
     GetResult,
     WhereDocument,
 )
+from chromadb.config import Component
 
-from chromadb.telemetry import Telemetry
-
-
-class API(ABC):
-    @abstractmethod
-    def __init__(self, telemetry_client: Telemetry):
-        pass
 
+class API(Component, ABC):
     @abstractmethod
     def heartbeat(self) -> int:
         """Returns the current server time in nanoseconds to check if the server is alive

chromadb/api/fastapi.py

@@ -7,7 +7,6 @@
     EmbeddingFunction,
     IDs,
     Include,
-    Metadata,
     Metadatas,
     Where,
     WhereDocument,
@@ -22,22 +21,27 @@
 from chromadb.api.models.Collection import Collection
 import chromadb.errors as errors
 from uuid import UUID
+from chromadb.telemetry import Telemetry
+from overrides import override
 
 
 class FastAPI(API):
     def __init__(self, system: System):
+        super().__init__(system)
         url_prefix = "https" if system.settings.chroma_server_ssl_enabled else "http"
         system.settings.require("chroma_server_host")
         system.settings.require("chroma_server_http_port")
         self._api_url = f"{url_prefix}://{system.settings.chroma_server_host}:{system.settings.chroma_server_http_port}/api/v1"
-        self._telemetry_client = system.get_telemetry()
+        self._telemetry_client = self.require(Telemetry)
 
+    @override
     def heartbeat(self) -> int:
         """Returns the current server time in nanoseconds to check if the server is alive"""
         resp = requests.get(self._api_url)
         raise_chroma_error(resp)
         return int(resp.json()["nanosecond heartbeat"])
 
+    @override
     def list_collections(self) -> Sequence[Collection]:
         """Returns a list of all collections"""
         resp = requests.get(self._api_url + "/collections")
@@ -49,6 +53,7 @@ def list_collections(self) -> Sequence[Collection]:
 
         return collections
 
+    @override
     def create_collection(
         self,
         name: str,
@@ -73,6 +78,7 @@ def create_collection(
             metadata=resp_json["metadata"],
         )
 
+    @override
     def get_collection(
         self,
         name: str,
@@ -90,6 +96,7 @@ def get_collection(
             metadata=resp_json["metadata"],
         )
 
+    @override
     def get_or_create_collection(
         self,
         name: str,
@@ -102,6 +109,7 @@ def get_or_create_collection(
             name, metadata, embedding_function, get_or_create=True
         )
 
+    @override
     def _modify(
         self,
         id: UUID,
@@ -115,11 +123,13 @@ def _modify(
         )
         raise_chroma_error(resp)
 
+    @override
     def delete_collection(self, name: str) -> None:
         """Deletes a collection"""
         resp = requests.delete(self._api_url + "/collections/" + name)
         raise_chroma_error(resp)
 
+    @override
     def _count(self, collection_id: UUID) -> int:
         """Returns the number of embeddings in the database"""
         resp = requests.get(
@@ -128,13 +138,15 @@ def _count(self, collection_id: UUID) -> int:
         raise_chroma_error(resp)
         return cast(int, resp.json())
 
-    def _peek(self, collection_id: UUID, limit: int = 10) -> GetResult:
+    @override
+    def _peek(self, collection_id: UUID, n: int = 10) -> GetResult:
         return self._get(
             collection_id,
-            limit=limit,
+            limit=n,
             include=["embeddings", "documents", "metadatas"],
         )
 
+    @override
     def _get(
         self,
         collection_id: UUID,
@@ -177,6 +189,7 @@ def _get(
             documents=body.get("documents", None),
         )
 
+    @override
     def _delete(
         self,
         collection_id: UUID,
@@ -196,6 +209,7 @@ def _delete(
         raise_chroma_error(resp)
         return cast(IDs, resp.json())
 
+    @override
     def _add(
         self,
         ids: IDs,
@@ -227,6 +241,7 @@ def _add(
         raise_chroma_error(resp)
         return True
 
+    @override
     def _update(
         self,
         collection_id: UUID,
@@ -255,6 +270,7 @@ def _update(
         resp.raise_for_status()
         return True
 
+    @override
     def _upsert(
         self,
         collection_id: UUID,
@@ -285,6 +301,7 @@ def _upsert(
         resp.raise_for_status()
         return True
 
+    @override
     def _query(
         self,
         collection_id: UUID,
@@ -320,18 +337,21 @@ def _query(
             documents=body.get("documents", None),
         )
 
+    @override
     def reset(self) -> bool:
         """Resets the database"""
         resp = requests.post(self._api_url + "/reset")
         raise_chroma_error(resp)
         return cast(bool, resp.json())
 
+    @override
     def persist(self) -> bool:
         """Persists the database"""
         resp = requests.post(self._api_url + "/persist")
         raise_chroma_error(resp)
         return cast(bool, resp.json())
 
+    @override
     def raw_sql(self, sql: str) -> pd.DataFrame:
         """Runs a raw SQL query against the database"""
         resp = requests.post(
@@ -340,6 +360,7 @@ def raw_sql(self, sql: str) -> pd.DataFrame:
         raise_chroma_error(resp)
         return pd.DataFrame.from_dict(resp.json())
 
+    @override
     def create_index(self, collection_name: str) -> bool:
         """Creates an index for the given space key"""
         resp = requests.post(
@@ -348,6 +369,7 @@ def create_index(self, collection_name: str) -> bool:
         raise_chroma_error(resp)
         return cast(bool, resp.json())
 
+    @override
     def get_version(self) -> str:
         """Returns the version of the server"""
         resp = requests.get(self._api_url + "/version")