[ENH]: Chroma client upgrade check

chromadb/__init__.py

@@ -37,6 +37,7 @@
     "UpdateCollectionMetadata",
     "QueryResult",
     "GetResult",
+    "__version__",
 ]
 
 logger = logging.getLogger(__name__)

chromadb/api/client.py

@@ -28,6 +28,7 @@
 from chromadb.telemetry.product.events import ClientStartEvent
 from chromadb.types import Database, Tenant, Where, WhereDocument
 import chromadb.utils.embedding_functions as ef
+from chromadb.utils.client_utils import _upgrade_check
 
 
 class SharedSystemClient:
@@ -137,6 +138,7 @@ def __init__(
         settings: Settings = Settings(),
     ) -> None:
         super().__init__(settings=settings)
+        _upgrade_check()
         self.tenant = tenant
         self.database = database
         # Create an admin client for verifying that databases and tenants exist

chromadb/app.py

@@ -1,7 +1,8 @@
-import chromadb
 import chromadb.config
 from chromadb.server.fastapi import FastAPI
+from chromadb.utils.client_utils import _upgrade_check
 
+_upgrade_check()
 settings = chromadb.config.Settings()
 server = FastAPI(settings)
 app = server.app()

chromadb/cli/cli.py

@@ -8,6 +8,7 @@
 import os
 import webbrowser
 
+from chromadb.utils.client_utils import _upgrade_check
 from chromadb.cli.utils import set_log_file_path
 
 app = typer.Typer()
@@ -58,6 +59,11 @@ def run(
         "\033[1mGetting started guide\033[0m: https://docs.trychroma.com/getting-started\n\n"
     )
 
+    upgrade_message = _upgrade_check()
+    if upgrade_message:
+        for m in upgrade_message:
+            typer.echo(m)
+
     # set ENV variable for PERSIST_DIRECTORY to path
     os.environ["IS_PERSISTENT"] = "True"
     os.environ["PERSIST_DIRECTORY"] = path
@@ -77,7 +83,6 @@ def run(
         "log_config": log_config,  # Pass the modified log_config dictionary
         "timeout_keep_alive": 30,
     }
-
     if test:
         return
 

chromadb/test/client/test_client_upgrade.py

@@ -0,0 +1,70 @@
+import json
+from unittest.mock import patch
+
+import pytest
+from pytest_httpserver import HTTPServer
+
+import chromadb
+
+
+def test_new_release_available(caplog: pytest.LogCaptureFixture) -> None:
+    with patch(
+        "chromadb.utils.client_utils._upgrade_check_url",
+        new="http://localhost:8008/pypi/chromadb/json",
+    ):
+        with HTTPServer(port=8008) as httpserver:
+            # Define the response
+            httpserver.expect_request("/pypi/chromadb/json").respond_with_data(
+                json.dumps({"info": {"version": "99.99.99"}})
+            )
+
+            # Your code that makes the HTTP call
+            chromadb.Client()
+
+            assert "A new release of chromadb is available" in caplog.text
+
+
+def test_on_latest_release(caplog: pytest.LogCaptureFixture) -> None:
+    with HTTPServer(port=8008) as httpserver:
+        # Define the response
+        httpserver.expect_request("/pypi/chromadb/json").respond_with_data(
+            json.dumps({"info": {"version": chromadb.__version__}})
+        )
+
+        # Your code that makes the HTTP call
+        chromadb.Client()
+
+        assert "A new release of chromadb is available" not in caplog.text
+
+
+def test_local_version_newer_than_latest(caplog: pytest.LogCaptureFixture) -> None:
+    with patch(
+        "chromadb.utils.client_utils._upgrade_check_url",
+        new="http://localhost:8008/pypi/chromadb/json",
+    ):
+        with HTTPServer(port=8008) as httpserver:
+            # Define the response
+            httpserver.expect_request("/pypi/chromadb/json").respond_with_data(
+                json.dumps({"info": {"version": "0.0.1"}})
+            )
+
+            # Your code that makes the HTTP call
+            chromadb.Client()
+
+            assert "A new release of chromadb is available" not in caplog.text
+
+
+def test_pypi_unavailable(caplog: pytest.LogCaptureFixture) -> None:
+    with patch(
+        "chromadb.utils.client_utils._upgrade_check_url",
+        new="http://localhost:8008/pypi/chromadb/json",
+    ):
+        with HTTPServer(port=8009) as httpserver:
+            # Define the response
+            httpserver.expect_request("/pypi/chromadb/json").respond_with_data(
+                json.dumps({"info": {"version": "99.99.99"}})
+            )
+
+            chromadb.Client()
+
+            assert "A new release of chromadb is available" not in caplog.text

chromadb/test/test_cli.py

@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 from typer.testing import CliRunner
 
 from chromadb.cli.cli import app
@@ -22,6 +24,27 @@ def test_app() -> None:
     assert "8001" in result.stdout
 
 
+def test_app_version_upgrade() -> None:
+    with patch(
+        "chromadb.__version__",
+        new="0.0.1",
+    ):
+        result = runner.invoke(
+            app,
+            [
+                "run",
+                "--path",
+                "chroma_test_data",
+                "--port",
+                "8001",
+                "--test",
+            ],
+        )
+    assert "A new release of chromadb is available" in result.stdout
+    assert "chroma_test_data" in result.stdout
+    assert "8001" in result.stdout
+
+
 def test_utils_set_log_file_path() -> None:
     log_config = set_log_file_path("chromadb/log_config.yml", "test.log")
     assert log_config["handlers"]["file"]["filename"] == "test.log"

chromadb/utils/client_utils.py

@@ -0,0 +1,62 @@
+import json
+import logging
+from typing import List
+from urllib import request
+
+logger = logging.getLogger(__name__)
+
+
+def compare_versions(version1: str, version2: str) -> int:
+    """Compares two versions of the format X.Y.Z and returns 1 if version1 is greater than version2, -1 if version1 is
+    less than version2, and 0 if version1 is equal to version2.
+    """
+    v1_components = list(map(int, version1.split(".")))
+    v2_components = list(map(int, version2.split(".")))
+
+    for v1, v2 in zip(v1_components, v2_components):
+        if v1 > v2:
+            return 1
+        elif v1 < v2:
+            return -1
+
+    if len(v1_components) > len(v2_components):
+        return 1
+    elif len(v1_components) < len(v2_components):
+        return -1
+
+    return 0
+
+
+_upgrade_check_url: str = "https://pypi.org/pypi/chromadb/json"
+_check_performed: bool = False
+
+
+def _upgrade_check() -> List[str]:
+    """Check pypi index for new version if possible."""
+    global _check_performed
+    upgrade_messages: List[str] = []
+    # this is to prevent cli from double printing
+    if _check_performed:
+        return upgrade_messages
+    try:
+        data = json.load(
+            request.urlopen(request.Request(_upgrade_check_url), timeout=5)
+        )
+        from chromadb import __version__ as local_chroma_version
+
+        latest_version = data["info"]["version"]
+        if compare_versions(latest_version, local_chroma_version) > 0:
+            upgrade_messages.append(
+                f"\033[38;5;069m[notice]\033[0m A new release of chromadb is available: "
+                f"\033[38;5;196m{local_chroma_version}\033[0m -> "
+                f"\033[38;5;082m{latest_version}\033[0m"
+            )
+            upgrade_messages.append(
+                "\033[38;5;069m[notice]\033[0m To upgrade, run `pip install --upgrade chromadb`."
+            )
+    except Exception:
+        pass
+    _check_performed = True
+    for m in upgrade_messages:
+        logger.info(m)
+    return upgrade_messages

requirements_dev.txt

@@ -8,6 +8,7 @@ mypy-protobuf
 pre-commit
 pytest
 pytest-asyncio
+pytest-httpserver
 setuptools_scm
 types-protobuf
 types-requests==2.30.0.0