From a8ffef385673db41c626e5b7dfcc85ff76c87786 Mon Sep 17 00:00:00 2001 From: David <30951234+Davidyz@users.noreply.github.com> Date: Tue, 15 Jul 2025 06:48:58 +0100 Subject: [PATCH 1/2] Revert "fix(cli): Verify chromadb connection by checking openapi title" This reverts commit bb52bcd79578c32e154b6f37968ca4a2c1cc635d. --- src/vectorcode/common.py | 25 ++++++-------- tests/test_common.py | 70 ++++++++++++++++++++++++++++++---------- 2 files changed, 63 insertions(+), 32 deletions(-) diff --git a/src/vectorcode/common.py b/src/vectorcode/common.py index 0c0cc536..c5f7cee4 100644 --- a/src/vectorcode/common.py +++ b/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 diff --git a/tests/test_common.py b/tests/test_common.py index c58b581d..c0dbdc5f 100644 --- a/tests/test_common.py +++ b/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 = ( From e424c2de0d9bcd5276f007d56323322dada9a99d Mon Sep 17 00:00:00 2001 From: Zhe Yu Date: Tue, 15 Jul 2025 16:09:53 +0800 Subject: [PATCH 2/2] feat(cli): Improve error message for chromadb connection issues --- src/vectorcode/main.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/vectorcode/main.py b/src/vectorcode/main.py index d56b90f6..fd14ae14 100644 --- a/src/vectorcode/main.py +++ b/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()