diff --git a/chromadb/__init__.py b/chromadb/__init__.py index 599fb94dd45..f8bc1617511 100644 --- a/chromadb/__init__.py +++ b/chromadb/__init__.py @@ -18,6 +18,8 @@ WhereDocument, UpdateCollectionMetadata, ) +from chromadb.utils import once + # Re-export types from chromadb.types __all__ = [ @@ -94,7 +96,7 @@ def configure(**kwargs) -> None: # type: ignore def get_settings() -> Settings: return __settings - +@once def EphemeralClient(settings: Settings = Settings()) -> API: """ Creates an in-memory instance of Chroma. This is useful for testing and @@ -104,7 +106,7 @@ def EphemeralClient(settings: Settings = Settings()) -> API: return Client(settings) - +@once def PersistentClient(path: str = "./chroma", settings: Settings = Settings()) -> API: """ Creates a persistent instance of Chroma that saves to disk. This is useful for diff --git a/chromadb/test/test_client.py b/chromadb/test/test_client.py index 1164e1e699d..070fdc167f1 100644 --- a/chromadb/test/test_client.py +++ b/chromadb/test/test_client.py @@ -3,6 +3,7 @@ import chromadb.server.fastapi import pytest import tempfile +import gc @pytest.fixture @@ -35,3 +36,33 @@ def test_persistent_client(persistent_api: API) -> None: def test_http_client(http_api: API) -> None: settings = http_api.get_settings() assert settings.chroma_api_impl == "chromadb.api.fastapi.FastAPI" + +def test_multiple_persistent_client() -> None: + with pytest.raises(RuntimeError): + a = chromadb.PersistentClient( + path=tempfile.gettempdir() + "/test_server", + ) + b = chromadb.PersistentClient( + path=tempfile.gettempdir() + "/test_server", + ) + +def test_multiple_ephemeral_client() -> None: + with pytest.raises(RuntimeError): + a = chromadb.EphemeralClient() + b = chromadb.EphemeralClient() + +def test_gc_ephemeral_client() -> None: + # need to manually gargabe collect otherwise tests fail from above instantiations. + gc.collect() + a = chromadb.EphemeralClient() + del a + gc.collect() + b = chromadb.EphemeralClient() + +def test_gc_persistent_client() -> None: + # need to manually gargabe collect otherwise tests fail from above instantiations. + gc.collect() + a = chromadb.PersistentClient(path=tempfile.gettempdir() + "/test_server",) + del a + gc.collect() + b = chromadb.PersistentClient(path=tempfile.gettempdir() + "/test_server",) diff --git a/chromadb/utils/__init__.py b/chromadb/utils/__init__.py index fe6bb81853b..1abfa2ea18a 100644 --- a/chromadb/utils/__init__.py +++ b/chromadb/utils/__init__.py @@ -1,4 +1,5 @@ import importlib +import weakref from typing import Type, TypeVar, cast C = TypeVar("C") @@ -10,3 +11,22 @@ def get_class(fqn: str, type: Type[C]) -> Type[C]: module = importlib.import_module(module_name) cls = getattr(module, class_name) return cast(Type[C], cls) + +def once(func): + """ + Decorator that limits a function to one call. + """ + instance_ref = None + def wrapper(*args, **kwargs): + nonlocal instance_ref + # if first run or the object created has been garbage collected. + if instance_ref is None or instance_ref() is None: + instance = func(*args, **kwargs) + instance_ref = weakref.ref(instance) + return instance + else: + raise RuntimeError(f"Function {func.__name__} has already been called. \ + You should try and use only one instance of the Client.") + return None + + return wrapper \ No newline at end of file