revert(cli): Revert "Verify chromadb connection by checking openapi title." (#246)

* Revert "fix(cli): Verify chromadb connection by checking openapi title"

This reverts commit bb52bcd.

* feat(cli): Improve error message for chromadb connection issues

Author: Davidyz

src/vectorcode/common.py

@@ -1,13 +1,11 @@
 import asyncio
 import contextlib
 import hashlib
-import json
 import logging
 import os
 import socket
 import subprocess
 import sys
-import traceback
 from asyncio.subprocess import Process
 from dataclasses import dataclass
 from typing import Any, AsyncGenerator, Optional
@@ -48,19 +46,16 @@ async def get_collections(
 
 
 async def try_server(base_url: str):
-    openapi_url = f"{base_url}/openapi.json"
-    try:
-        async with httpx.AsyncClient() as client:
-            response = await client.get(url=openapi_url)
-            logger.debug(f"Fetching openapi.json from {openapi_url}: {response=}")
-            if response.status_code != 200:
-                return False
-            openapi_json = json.loads(response.content.decode())
-            if openapi_json:
-                return openapi_json.get("info", {}).get("title", "").lower() == "chroma"
-    except Exception as e:
-        logger.info(f"Failed to connect to chromadb at {base_url}")
-        logger.debug(traceback.format_exception(e))
+    for ver in ("v1", "v2"):  # v1 for legacy, v2 for latest chromadb.
+        heartbeat_url = f"{base_url}/api/{ver}/heartbeat"
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(url=heartbeat_url)
+                logger.debug(f"Heartbeat {heartbeat_url} returned {response=}")
+                if response.status_code == 200:
+                    return True
+        except (httpx.ConnectError, httpx.ConnectTimeout):
+            pass
     return False
 
 

src/vectorcode/main.py

@@ -4,6 +4,8 @@
 import sys
 import traceback
 
+import httpx
+
 from vectorcode import __version__
 from vectorcode.cli_utils import (
     CliAction,
@@ -100,8 +102,12 @@ async def async_main():
                 from vectorcode.subcommands import files
 
                 return_val = await files(final_configs)
-    except Exception:
+    except Exception as e:
         return_val = 1
+        if isinstance(e, httpx.RemoteProtocolError):  # pragma: nocover
+            e.add_note(
+                f"Please verify that {final_configs.db_url} is a working chromadb server."
+            )
         logger.error(traceback.format_exc())
     finally:
         await ClientManager().kill_servers()

tests/test_common.py

@@ -97,6 +97,59 @@ def test_get_embedding_function_init_exception():
         )
 
 
+@pytest.mark.asyncio
+async def test_try_server_versions():
+    # Test successful v1 response
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_client.return_value.__aenter__.return_value.get.return_value = (
+            mock_response
+        )
+        assert await try_server("http://localhost:8300") is True
+        mock_client.return_value.__aenter__.return_value.get.assert_called_once_with(
+            url="http://localhost:8300/api/v1/heartbeat"
+        )
+
+    # Test fallback to v2 when v1 fails
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_response_v1 = MagicMock()
+        mock_response_v1.status_code = 404
+        mock_response_v2 = MagicMock()
+        mock_response_v2.status_code = 200
+        mock_client.return_value.__aenter__.return_value.get.side_effect = [
+            mock_response_v1,
+            mock_response_v2,
+        ]
+        assert await try_server("http://localhost:8300") is True
+        assert mock_client.return_value.__aenter__.return_value.get.call_count == 2
+
+    # Test both versions fail
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_response_v1 = MagicMock()
+        mock_response_v1.status_code = 404
+        mock_response_v2 = MagicMock()
+        mock_response_v2.status_code = 500
+        mock_client.return_value.__aenter__.return_value.get.side_effect = [
+            mock_response_v1,
+            mock_response_v2,
+        ]
+        assert await try_server("http://localhost:8300") is False
+
+    # Test connection error cases
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_client.return_value.__aenter__.return_value.get.side_effect = (
+            httpx.ConnectError("Cannot connect")
+        )
+        assert await try_server("http://localhost:8300") is False
+
+    with patch("httpx.AsyncClient") as mock_client:
+        mock_client.return_value.__aenter__.return_value.get.side_effect = (
+            httpx.ConnectTimeout("Connection timeout")
+        )
+        assert await try_server("http://localhost:8300") is False
+
+
 def test_verify_ef():
     # Mocking AsyncCollection and Config
     mock_collection = MagicMock()
@@ -137,18 +190,10 @@ async def test_try_server_mocked(mock_socket):
     with patch("httpx.AsyncClient") as mock_client:
         mock_response = MagicMock()
         mock_response.status_code = 200
-        mock_response.content = b'{"info":{"title": "Chroma"}}'
         mock_client.return_value.__aenter__.return_value.get.return_value = (
             mock_response
         )
         assert await try_server("http://localhost:8000") is True
-    with patch("httpx.AsyncClient") as mock_client:
-        mock_response = MagicMock()
-        mock_response.status_code = 404
-        mock_client.return_value.__aenter__.return_value.get.return_value = (
-            mock_response
-        )
-        assert await try_server("http://localhost:8000") is False
 
     # Mocking httpx.AsyncClient to raise a ConnectError
     with patch("httpx.AsyncClient") as mock_client:
@@ -157,15 +202,6 @@ async def test_try_server_mocked(mock_socket):
         )
         assert await try_server("http://localhost:8000") is False
 
-    with patch("httpx.AsyncClient") as mock_client:
-        mock_response = MagicMock()
-        mock_response.status_code = 200
-        mock_response.content = b'{"info":{"title": "Dummy"}}'
-        mock_client.return_value.__aenter__.return_value.get.return_value = (
-            mock_response
-        )
-        assert await try_server("http://localhost:8000") is False
-
     # Mocking httpx.AsyncClient to raise a ConnectTimeout
     with patch("httpx.AsyncClient") as mock_client:
         mock_client.return_value.__aenter__.return_value.get.side_effect = (